diff --git a/.gitattributes b/.gitattributes index 264fdfd02f..75fba5d6ab 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,2 +1,7 @@ # Normalize Python files to LF line endings *.py text eol=lf + +# Always check out shell scripts with LF endings. Without this rule a Windows +# clone (core.autocrlf=true) rewrites them to CRLF, and the trailing \r breaks +# them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -"). +*.sh text eol=lf diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2802f461b6..a9694036cf 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -6,10 +6,10 @@ /unsloth/models/rl_replacements.py @Datta0 @pluesclues @danielhanchen /unsloth/trainer.py @danielhanchen /unsloth/models/sentence_transformer.py @Etherll @danielhanchen -/unsloth/save.py @rolandtannous @danielhanchen +/unsloth/save.py @danielhanchen /unsloth/tokenizer_utils.py @mmathew23 @danielhanchen -/unsloth/chat_templates.py @rolandtannous @danielhanchen -/unsloth/ollama_template_mappers.py @rolandtannous @danielhanchen +/unsloth/chat_templates.py @danielhanchen +/unsloth/ollama_template_mappers.py @danielhanchen /unsloth/kernels/moe/*.py @Datta0 /unsloth/import_fixes.py @danielhanchen /unsloth/device_type.py @danielhanchen @@ -45,14 +45,14 @@ /unsloth/utils/hf_hub.py @mmathew23 /unsloth/utils/packing.py @mmathew23 -/cli/ @rolandtannous @Manan17 -/studio/frontend/ @Shine1i @rolandtannous @Manan17 +/cli/ @Manan17 +/studio/frontend/ @Shine1i @Manan17 /studio/frontend/public/ @Shine1i -/studio/backend/ @rolandtannous -/studio/backend/core/data_recipe/ @rolandtannous -/studio/backend/tests/ @rolandtannous @danielhanchen -/tests/ @rolandtannous @danielhanchen -/scripts/ @rolandtannous @danielhanchen +/studio/backend/ +/studio/backend/core/data_recipe/ +/studio/backend/tests/ @danielhanchen +/tests/ @danielhanchen +/scripts/ @danielhanchen # Snapshot data for the notebook linter / Colab oracle. Drift in these # files changes the pin floor for every Unsloth notebook, so refreshes diff --git a/.github/scripts/assert-llama-loads.sh b/.github/scripts/assert-llama-loads.sh new file mode 100755 index 0000000000..c2ffe27469 --- /dev/null +++ b/.github/scripts/assert-llama-loads.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Assert Studio installed a llama.cpp that loads and runs on THIS macOS. Tests +# the contract that matters (binaries load and their minimum-OS is <= this host) +# instead of the old "did install.sh fall back to a source build?" grep, since a +# source build with a correct deployment target is a valid outcome. +set -uo pipefail + +UNSLOTH_HOME="${STUDIO_HOME:-$HOME/.unsloth}" +LLAMA_DIR="${LLAMA_CPP_DIR:-$UNSLOTH_HOME/llama.cpp}" +BIN_DIR="$LLAMA_DIR/build/bin" + +fail() { + echo "::error::$*" + if [ -f logs/install.log ]; then + echo "---- install.log (llama.cpp lines) ----" + grep -E "llama-prebuilt|llama\.cpp|macos prebuilt|falling back" logs/install.log | tail -80 || true + fi + exit 1 +} + +SERVER="$(find "$LLAMA_DIR" -type f -name 'llama-server' 2>/dev/null | head -1)" +QUANT="$(find "$LLAMA_DIR" -type f -name 'llama-quantize' 2>/dev/null | head -1)" +[ -n "$SERVER" ] || fail "llama-server not found under $LLAMA_DIR after install" +[ -n "$QUANT" ] || fail "llama-quantize not found under $LLAMA_DIR after install" + +HOST_VER="$(sw_vers -productVersion 2>/dev/null || echo '0')" +HOST_MAJOR="${HOST_VER%%.*}" + +# Static minimum-OS check on every Mach-O we ship. vtool ships with the Xcode +# command line tools, which GitHub macOS runners always have; if it is somehow +# missing we skip the static check and rely on the runtime launch below. +if command -v vtool >/dev/null 2>&1; then + while IFS= read -r macho; do + [ -n "$macho" ] || continue + minos="$(vtool -show-build "$macho" 2>/dev/null | awk '/minos/{print $2; exit}')" + [ -n "$minos" ] || continue + min_major="${minos%%.*}" + if [ "$min_major" -gt "$HOST_MAJOR" ] 2>/dev/null; then + fail "$(basename "$macho") is built for macOS $minos but this runner is macOS $HOST_VER (prebuilt is newer than the host)" + fi + done < <(find "$BIN_DIR" -type f \( -name '*.dylib' -o -name 'llama-server' -o -name 'llama-quantize' \) 2>/dev/null) +fi + +# Runtime launch: --version forces dyld to load every linked dylib (including +# libggml-metal.dylib). A missing Metal symbol or too-new binary fails here. +if ! "$SERVER" --version >/tmp/llama-server-version.txt 2>&1; then + echo "---- llama-server --version output ----" + cat /tmp/llama-server-version.txt || true + fail "llama-server failed to launch on macOS $HOST_VER (dyld load / symbol error)" +fi + +echo "llama.cpp load validation passed on macOS $HOST_VER" +echo " server: $SERVER" +sed -n '1,4p' /tmp/llama-server-version.txt 2>/dev/null || true diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index f8eccc7d18..f7c338d76b 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -269,7 +269,8 @@ jobs: tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ tests/utils/test_attention_masks.py \ - tests/utils/test_trunc_normal_patch.py + tests/utils/test_trunc_normal_patch.py \ + tests/python/test_fast_language_model_text_only.py python -m pytest --collect-only -q "$RUNNER_TEMP/unsloth-zoo/tests/" - name: import_fixes drift detectors (18 tests, HARD GATE) @@ -332,12 +333,21 @@ jobs: run: | python -m pytest -v --tb=short tests/test_callback_signature_drift.py + - name: generation correctness guards (HARD GATE) + # Deterministic CPU guards, each validated to fail on its pre-fix code: + # leftpad = batched left-padded generation (#1066/#3699, fixed by + # #2216 + #4100; staging proof: unsloth-staging-2 PRs 170/172); + # rope_scaling_drift = config.rope_scaling dropped by replaced rotary + # classes (#2405). AST checks run first so import breakage cannot mask them. + run: | + python -m pytest -v --tb=short \ + tests/utils/test_prepare_inputs_leftpad.py \ + tests/utils/test_rope_scaling_drift.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 - # because their sibling files need real GPUs / real HF weights. - # The five files below are pure-Python + AST/protobuf/regex tests - # that run cleanly on CPU. Env inherited from the job block. + # CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/ + # that Repo tests (CPU) --ignores. AST/protobuf/regex plus tiny CPU model + # loads; run cleanly here (transformers/torch installed). run: | python -m pytest -q --tb=short \ tests/saving/test_save_shell_injection.py \ @@ -345,11 +355,12 @@ jobs: tests/saving/test_fix_sentencepiece_gguf_robustness.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ + tests/python/test_fast_language_model_text_only.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn # requires CUDA + dev toolchain, which the CPU-only ubuntu-latest - # runner does not have. The other 15 Bucket-A tests pass cleanly. + # runner does not have. The other Bucket-A tests pass cleanly. - name: unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} — full pytest (CPU) # 106 of 111 test_* in unsloth_zoo are CPU-only. The two CUDA-skip @@ -990,8 +1001,10 @@ jobs: # 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", + "deepseek_ocr2": "TimeoutError: compile exceeds per-model budget", "sam": "TimeoutError: compile exceeds per-model budget", "sam_hq": "TimeoutError: compile exceeds per-model budget", + "deepseek_ocr2": "TimeoutError: compile exceeds per-model budget", } diff --git a/.github/workflows/lint-ci.yml b/.github/workflows/lint-ci.yml index 00e6e357e2..bd859a6e9e 100644 --- a/.github/workflows/lint-ci.yml +++ b/.github/workflows/lint-ci.yml @@ -79,6 +79,64 @@ jobs: run: | ruff check unsloth unsloth_cli studio tests cli.py unsloth-cli.py + - name: Import-hoist verifier self-test + # scripts/verify_import_hoist.py is a scope-aware (LEGB) AST + # resolver that gates import-hoisting / alias-rename refactors + # against two bugs ruff and pyflakes both miss: + # 1. dangling alias -- `from a import b as _b` hoisted to + # `from a import b` but a leftover `_b` reference now + # resolves to nothing (or to some other module-level `_b`). + # 2. rename clash -- `_b -> b` silently re-points at a + # different object already named `b` in that scope. + # This step runs the tool's 8 negative-control cases so a + # regression in the verifier itself fails before we trust it on + # a diff. Hermetic, stdlib-only, sub-second. Hard gate. + run: | + python scripts/verify_import_hoist.py --self-test + + - name: Import-hoist / alias-rename safety (changed Python files) + # Runs the verifier in compare mode on every in-place-modified + # .py in the PR: parses each file BEFORE (base branch) and AFTER + # (this diff), resolves every name load, and fails on a BLOCKER + # (dangling alias / rename clash / re-pointed import). INFO + # findings (a helper relocated to another file) do not fail. + # + # --diff-filter=M (in-place edits only) is deliberate: that is + # exactly where a hoist refactor lives, and it skips brand-new + # files whose re-export imports would otherwise look "unused". + # + # Diff against the true merge-base, not the base tip. A two-dot + # diff against the tip re-lints every file the base branch + # changed after the PR branched, comparing newer base code + # (BEFORE) against the PR's older snapshot (AFTER) - a + # time-reversed comparison that flags the base branch's own + # refactors as blockers on PRs that never touched those files. + # The compare API returns the merge-base without needing local + # history, and fetching that single commit by SHA keeps the + # shallow (fetch-depth: 1) clone. + if: github.event_name == 'pull_request' + env: + GH_TOKEN: ${{ github.token }} + run: | + MERGE_BASE=$(gh api \ + "repos/${{ github.repository }}/compare/${{ github.event.pull_request.base.sha }}...${{ github.event.pull_request.head.sha }}" \ + --jq .merge_base_commit.sha) + git fetch --no-tags --depth=1 origin "$MERGE_BASE" + mapfile -t CHANGED < <( + git diff --name-only --diff-filter=M \ + "$MERGE_BASE" HEAD -- '*.py' \ + | grep -vE '(^|/)(unsloth_compiled_cache|node_modules|build|dist)/' || true + ) + if [ "${#CHANGED[@]}" -eq 0 ]; then + echo "no in-place-modified Python files to check" + exit 0 + fi + printf 'merge base: %s\n' "$MERGE_BASE" + printf 'checking %d file(s):\n' "${#CHANGED[@]}" + printf ' %s\n' "${CHANGED[@]}" + python scripts/verify_import_hoist.py \ + --before "$MERGE_BASE" --after HEAD "${CHANGED[@]}" + - name: No leftover debugger / pdb / breakpoint calls # Catches the "I'll just stick a breakpoint() here" mistake # before it ships. AST-based so commented-out debugger diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index 75940832a0..221e86f235 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -231,33 +231,14 @@ jobs: tests/studio/test_is_mlx_dispatch_gate.py \ tests/studio/test_mlx_training_worker_behaviors.py - # Studio prebuilt llama.cpp install + GGUF inference. Drives the - # exact path Studio's setup.sh takes on macOS: invokes - # studio/install_llama_prebuilt.py with --published-repo - # ggml-org/llama.cpp and --published-release-tag b9049 (the - # latest llama.cpp release at the time this step was added; bump - # via UNSLOTH_LLAMA_TAG / DEFAULT_LLAMA_TAG when refreshing). - # The installer downloads llama-b9049-bin-macos-arm64.tar.gz, - # which is the universal Apple Silicon (arm64) build -- the - # same artifact works on M1/M2/M3/M4 because llama.cpp compiles - # against the ARMv8.2 baseline. - # - # The b9049 release also publishes: - # - llama-b9049-bin-macos-arm64-kleidiai.tar.gz - # KleidiAI dispatches at runtime; on M1 it falls back where - # ISA features (e.g. I8MM) are missing, so this asset also - # runs on M1 -- Studio just doesn't choose it by default. - # - llama-b9049-bin-macos-x64.tar.gz - # Intel-only; would only run on M1 via Rosetta 2 emulation, - # which we explicitly avoid. - # - iOS XCFramework - # iOS-app build artifact, unrelated to a macOS desktop CI. - # - # After install, downloads a small published GGUF - # (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) from HuggingFace and - # runs the prebuilt llama-cli on it. Asserts the prompt echo - # appears in stdout. If the install fails OR the binary exits - # non-zero, that's an Unsloth/Studio bug. + # Studio prebuilt llama.cpp install + GGUF inference. Mirrors the + # path Studio's setup.sh takes on macOS since #5963: plan against + # the unslothai/llama.cpp fork's latest release, which ships the + # bin-macos-arm64 bundle plus the llama-prebuilt-manifest.json the + # default policy reads. After install, downloads a small published + # GGUF (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) and validates + # llama-server /completion end to end. An install failure or a + # non-zero binary exit is an Unsloth/Studio bug. - name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1) env: HF_TOKEN: ${{ secrets.HF_TOKEN }} @@ -272,20 +253,12 @@ jobs: set -euo pipefail INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" rm -rf "$INSTALL_DIR" - # --simple-policy is required when --published-repo points - # at upstream ggml-org/llama.cpp; that repo doesn't ship the - # llama-prebuilt-manifest.json asset Studio's default policy - # expects, so the simple platform-specific policy maps - # Darwin+arm64 -> bin-macos-arm64 directly. studio/setup.sh - # passes both --published-repo ggml-org/llama.cpp AND - # --simple-policy automatically on macOS, so this CI step - # exercises the same code path users hit when they run - # `curl -fsSL https://unsloth.ai/install.sh | sh`. + # Mirror studio/setup.sh on macOS (the install.sh user path): + # it plans against the unslothai/llama.cpp fork's latest + # release with no policy or tag flags. python studio/install_llama_prebuilt.py \ --install-dir "$INSTALL_DIR" \ - --published-repo ggml-org/llama.cpp \ - --published-release-tag b9049 \ - --simple-policy + --published-repo unslothai/llama.cpp # Studio bundles only llama-server + llama-quantize from the # prebuilt (not llama-cli) -- inference goes through diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml index 673b2f3cc5..2edcae8ab2 100644 --- a/.github/workflows/notebooks-ci.yml +++ b/.github/workflows/notebooks-ci.yml @@ -285,7 +285,15 @@ jobs: # The PR-time CI must validate the code in this PR; PyPI unsloth # may lag the in-repo CPU-torch fallback in unsloth/kernels/utils.py # (lines 162-170) that handles missing torch._C._cuda_getCurrentRawStream. - pip install --no-deps unsloth_zoo + # unsloth_zoo from git main mirrors every other CI (Core / MLX / + # install.sh) so PR-time validation sees the same zoo HEAD. + for attempt in 1 2 3; do + if pip install --no-deps "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done pip install --no-deps -e ./unsloth - name: Convert notebooks for AST scan diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index a1e7b2efa6..33ac3b9bd8 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -72,6 +72,31 @@ concurrency: permissions: contents: read +# ────────────────────────────────────────────────────────────────────── +# Network-resilience knobs, applied to every job/step. These add retries +# and backoff ONLY; they do not relax a single integrity check. cargo +# still resolves against Cargo.lock (--locked), pip still verifies the +# wheels it downloads, npm still enforces package-lock integrity, the +# harden-runner egress allowlists below are unchanged, and every action +# stays SHA-pinned. The advisory-audit run on 2026-05-29 red-failed when +# one crates.io tarball fetch hit "Recv failure: Connection reset by +# peer" (curl 56); cargo's default of 3 retries over an HTTP/2-multiplexed +# connection did not recover. The settings below make that class of +# transient fault self-heal instead of failing the whole run. +env: + # pip: raise the built-in retry count and per-connection timeout. + PIP_RETRIES: "10" + PIP_DEFAULT_TIMEOUT: "60" + # cargo: retry network ops and disable HTTP/2 multiplexing -- the + # documented mitigation for the curl-56 connection resets above. + CARGO_NET_RETRY: "10" + CARGO_HTTP_MULTIPLEXING: "false" + CARGO_NET_GIT_FETCH_WITH_CLI: "true" + # npm: retry registry fetches with capped exponential backoff. + NPM_CONFIG_FETCH_RETRIES: "5" + NPM_CONFIG_FETCH_RETRY_MINTIMEOUT: "2000" + NPM_CONFIG_FETCH_RETRY_MAXTIMEOUT: "60000" + jobs: # ───────────────────────────────────────────────────────────────────── # Combined advisory-DB audit: pip-audit + npm audit + cargo audit @@ -140,7 +165,7 @@ jobs: - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27 - - uses: swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2.9.1 + - uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 with: workspaces: studio/src-tauri -> target @@ -153,8 +178,23 @@ jobs: # crashes with a TOML parse error on that file. # npm audit is bundled with the node toolchain, no install. run: | - python -m pip install --upgrade pip 'pip-audit>=2.7' - cargo install --locked --version '^0.22' cargo-audit + retry() { # retry with exponential backoff + local max="$1"; shift + local n=1 delay=5 + until "$@"; do + if [ "$n" -ge "$max" ]; then + echo "::error::command failed after ${n} attempts: $*" >&2 + return 1 + fi + echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2 + sleep "$delay"; n=$((n + 1)); delay=$((delay * 2)) + done + } + retry 5 python -m pip install --upgrade pip 'pip-audit>=2.7' + # --locked keeps the resolved tree identical to Cargo.lock; the + # CARGO_NET_* env above plus this outer loop survive transient + # crates.io connection resets without weakening that guarantee. + retry 5 cargo install --locked --version '^0.22' cargo-audit # ───────────────────────────────────────────────────────────── # Python: pip-audit @@ -330,32 +370,60 @@ jobs: # ───────────────────────────────────────────────────────────── # OSV-Scanner: cross-ecosystem advisory DB (PyPI + npm + cargo) # ───────────────────────────────────────────────────────────── + - name: Download + verify OSV-Scanner + # Split out from the scan below so binary integrity is a HARD gate: + # a checksum mismatch (swapped release asset, the Trivy-style pivot + # this workflow refuses) fails the job instead of being swallowed by + # the scan step's continue-on-error. A download still failing after + # retries is transient, so we skip the scan rather than red-fail. + # SHA-256 verified BEFORE chmod +x / exec. Bump OSV_SHA256 in lockstep + # with OSV_VERSION (value from the release's osv-scanner_SHA256SUMS). + run: | + set -euo pipefail + OSV_VERSION="v2.0.2" + OSV_SHA256="3abcfd7126c453a00421487e721b296e0cb68085bd431d6cef60872774170fc8" + if ! curl --proto '=https' --tlsv1.2 -fsSL \ + --retry 5 --retry-delay 3 --retry-connrefused --retry-all-errors \ + -o /tmp/osv-scanner \ + "https://github.com/google/osv-scanner/releases/download/${OSV_VERSION}/osv-scanner_linux_amd64"; then + echo "::warning::osv-scanner download failed after retries; skipping scan" >&2 + rm -f /tmp/osv-scanner + exit 0 # transient availability: do not red-fail the job + fi + if ! echo "${OSV_SHA256} /tmp/osv-scanner" | sha256sum -c -; then + echo "::error::osv-scanner checksum mismatch; refusing to execute" >&2 + rm -f /tmp/osv-scanner + exit 1 # integrity failure: hard-fail + fi + chmod +x /tmp/osv-scanner + /tmp/osv-scanner --version + - name: OSV-Scanner (PyPI + npm + cargo, cross-ecosystem advisories) # OSV's advisory feed is a superset of GitHub-Advisory + RustSec # + npm advisories; running it alongside the per-ecosystem audit # tools catches CVEs that haven't propagated to the per-ecosystem # DBs yet (e.g. langchain-core CVE-2025-68664 was on OSV before # GitHub Advisory). Single binary, one transitive resolver, all - # three lockfile types in one pass. Non-blocking until baselines - # close. + # three lockfile types in one pass. Binary is checksum-verified in + # the step above; only the advisory scan stays non-blocking until + # baselines close. continue-on-error: true run: | set +e - # OSV-Scanner ships a raw binary (no tarball) in v2.x. - curl -fsSL -o /tmp/osv-scanner \ - https://github.com/google/osv-scanner/releases/download/v2.0.2/osv-scanner_linux_amd64 - chmod +x /tmp/osv-scanner - /tmp/osv-scanner --version - /tmp/osv-scanner scan source \ - --lockfile=studio/frontend/package-lock.json \ - --lockfile=studio/src-tauri/Cargo.lock \ - --lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \ - --lockfile=requirements.txt:audit-reqs/studio.txt \ - --lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \ - --lockfile=requirements.txt:audit-reqs/overrides.txt \ - --lockfile=requirements.txt:audit-reqs/extras.txt \ - --lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \ - --format=table 2>&1 | tee logs-osv-scanner.txt + if [ ! -x /tmp/osv-scanner ]; then + echo "osv-scanner unavailable this run; skipping scan" | tee logs-osv-scanner.txt + else + /tmp/osv-scanner scan source \ + --lockfile=studio/frontend/package-lock.json \ + --lockfile=studio/src-tauri/Cargo.lock \ + --lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \ + --lockfile=requirements.txt:audit-reqs/studio.txt \ + --lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \ + --lockfile=requirements.txt:audit-reqs/overrides.txt \ + --lockfile=requirements.txt:audit-reqs/extras.txt \ + --lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \ + --format=table 2>&1 | tee logs-osv-scanner.txt + fi { echo "## OSV-Scanner (cross-ecosystem)" echo @@ -1075,7 +1143,23 @@ jobs: # new-install-script gate below protects against, and we must # not run any third-party hook to set up the audit. working-directory: studio/frontend - run: npm ci --ignore-scripts + run: | + retry() { # retry with exponential backoff + local max="$1"; shift + local n=1 delay=5 + until "$@"; do + if [ "$n" -ge "$max" ]; then + echo "::error::command failed after ${n} attempts: $*" >&2 + return 1 + fi + echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2 + sleep "$delay"; n=$((n + 1)); delay=$((delay * 2)) + done + } + # --ignore-scripts is mandatory here (no third-party hook runs); + # the retry only re-attempts the registry fetch, it never relaxes + # that flag or the package-lock integrity check npm ci enforces. + retry 5 npm ci --ignore-scripts - name: npm audit signatures (informational) # Surfaces unsigned / mis-signed packages from the npm diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index 53514e2ce1..b196805cf7 100644 --- a/.github/workflows/studio-api-smoke.yml +++ b/.github/workflows/studio-api-smoke.yml @@ -77,7 +77,7 @@ jobs: path: hf-cache # Same key as studio-ui-smoke.yml so the two jobs share a # single GGUF download across CI. - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -88,17 +88,19 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 63eb70f7f1..88c7344683 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -144,9 +144,19 @@ jobs: # versions ship a CPU build that imports cleanly on Linux. pip install 'bitsandbytes>=0.45' # unsloth.device_type imports unsloth_zoo.utils.Version at module - # scope, so the conftest preload needs unsloth_zoo even though - # it is an optional dep of unsloth. - pip install 'unsloth_zoo>=2026.5.1' + # scope, so the conftest preload needs unsloth_zoo. Pull from + # git main so this job sees the same zoo HEAD as Core / MLX / + # install.sh do (otherwise a fix on zoo main hides until release). + # No --no-deps: matches prior `pip install 'unsloth_zoo>=2026.5.1'` + # behaviour so triton etc. still come in for the Repo tests CPU + # collection imports. + for attempt in 1 2 3; do + if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done pip install -e . --no-deps - name: Repo tests (CPU, auto-discovered) @@ -212,6 +222,7 @@ jobs: for s in \ tests/sh/test_get_torch_index_url.sh \ tests/sh/test_mac_intel_compat.sh \ + tests/sh/test_nvcc_meets_llama_minimum.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh; do echo "::group::$s" diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml index 1270a57ef6..b42086f191 100644 --- a/.github/workflows/studio-frontend-ci.yml +++ b/.github/workflows/studio-frontend-ci.yml @@ -17,6 +17,8 @@ on: - 'studio/frontend/**' - 'scripts/check_frontend_dep_removal.py' - 'tests/studio/test_frontend_dep_removal.py' + - 'scripts/sync_allow_scripts_pins.py' + - 'tests/studio/test_sync_allow_scripts_pins.py' - '.github/workflows/studio-frontend-ci.yml' push: branches: [main, pip] @@ -60,6 +62,19 @@ jobs: with: node-version: '22' + # node 22 bundles npm 10.x, which predates allowScripts. Move to the + # 11.x line and fail loudly if the gate is still missing, so the + # strict flag below can never silently degrade into a warning. + - name: Upgrade npm to 11.x (allowScripts enforcement) + working-directory: ${{ github.workspace }} + run: | + npm install -g npm@^11 --no-fund --no-audit + V=$(npm -v) + case "$V" in + 11.1[6-9].*|11.[2-9][0-9].*|1[2-9].*) echo "npm $V has allowScripts" ;; + *) echo "::error::npm $V lacks allowScripts (need >=11.16)"; exit 1 ;; + esac + # Run the structural lockfile scan BEFORE npm ci. A compromised # tarball runs its `prepare` / `postinstall` during `npm ci`, # so any catch has to fire upstream of that. The scanner is @@ -68,14 +83,23 @@ jobs: working-directory: ${{ github.workspace }} run: python3 scripts/lockfile_supply_chain_audit.py + # Dependency bumps strand the version-pinned allowScripts entries. + # The paired pre-commit hook auto-fixes PRs; this is the backstop. + - name: allowScripts pins must match the lockfile + working-directory: ${{ github.workspace }} + run: | + python3 tests/studio/test_sync_allow_scripts_pins.py + python3 scripts/sync_allow_scripts_pins.py --check + - name: Lockfile must agree with package.json (npm ci is strict) - # Lifecycle scripts (esbuild native-binary postinstall, etc.) are - # required for `vite build`. The pre-install lockfile structural - # audit (lockfile_supply_chain_audit.py) is the practical defence - # against the npm postinstall-dropper class -- it fires BEFORE any - # tarball runs, on the injection pattern itself rather than an - # advisory-DB lookup. - run: npm ci --no-fund --no-audit + # The vite 8 chain (rolldown, lightningcss, tailwind oxide) ships napi + # binaries with no install scripts. The only script-bearing deps are + # covered by `allowScripts` in package.json (npm >=11.16, default in + # npm 12). The pre-install lockfile audit above stays the first line + # of defence -- it fires before any tarball can run code. + # --strict-allow-scripts: any unreviewed install script hard-fails + # the job; the sync hook keeps the pins fresh after bumps. + run: npm ci --strict-allow-scripts --no-fund --no-audit - name: npm ci must not have modified the working tree working-directory: ${{ github.workspace }} diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index 6def56f769..cffb33f71d 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -20,7 +20,7 @@ # enable_tools / enabled_tools, and enable_thinking on/off. # # 3. JSON, images -# gemma-4-E2B-it UD-IQ3_XXS (~2.4 GiB) + mmproj-F16 (~986 MiB). +# Qwen3-VL-2B-Instruct UD-Q4_K_XL (~1.1 GiB) + mmproj-F16 (~780 MiB). # response_format JSON-schema decoding and OpenAI image_url # (data URI) plus Anthropic source/base64 image inputs. # @@ -91,7 +91,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -102,17 +102,19 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -296,6 +298,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: openai-anthropic-log @@ -373,6 +377,7 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -771,6 +776,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: tool-calling-log @@ -787,9 +794,15 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 30 env: - GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF - GGUF_VARIANT: UD-IQ3_XXS - GGUF_FILE: gemma-4-E2B-it-UD-IQ3_XXS.gguf + GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF + # UD-Q4_K_XL, not UD-IQ2_XXS: at 2-bit the temp-0 answer to the JSON + # step's capital-of-France probe flips with the host's SIMD kernels + # (GitHub runners deterministically answered France while other CPUs + # answer Paris; seeds do not rescue it, 1/5 Paris at temp 0.7). The + # Q4 quant answered Paris 13/13 across temps and seeds on the same + # runners, so the hard Paris assertion below stays reliable. + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: Qwen3-VL-2B-Instruct-UD-Q4_K_XL.gguf MMPROJ_FILE: mmproj-F16.gguf STUDIO_PORT: '18890' HF_HOME: ${{ github.workspace }}/hf-cache @@ -819,7 +832,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - name: Prime HF_HOME with the GGUF + mmproj id: prime-hf @@ -831,17 +844,19 @@ jobs: mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} (model + mmproj) if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -884,13 +899,23 @@ jobs: -H 'content-type: application/json' \ -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" - # Load the GGUF (mmproj is auto-detected via the HF repo - # lookup, the cached file is pulled out of HF_HOME). - curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ - -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ - --max-time 900 \ - -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}" \ - | jq '{status, display_name, is_vision}' + # Retry: llama-server startup can race process teardown after a + # failed attempt. Keep curl out of a pipe so HTTP failures are not + # masked by jq. + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 900 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP; response:" + cat /tmp/load.json || true + sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name, is_vision}' /tmp/load.json - name: JSON schema decoding + image input env: @@ -1043,6 +1068,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: json-images-log diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index b4e274155e..412726538c 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -62,7 +62,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -73,29 +73,26 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Install pyjwt for the JWT-expiry forge test run: pip install 'pyjwt>=2.6' diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index fab0a36bd1..c794a34acd 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -85,7 +85,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -96,6 +96,7 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf # Save partial caches on cancel/timeout -- hf download resumes by # content hash. `outcome != skipped` keeps cache-hit a no-op. @@ -104,23 +105,19 @@ jobs: uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Install OpenAI + Anthropic Python SDKs run: pip install 'openai>=1.50' 'anthropic>=0.40' @@ -294,6 +291,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: openai-anthropic-log @@ -364,18 +363,14 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Reset auth + boot Studio (API-only, default tool policy) # We deliberately use the API-only mode rather than @@ -659,6 +654,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: tool-calling-log @@ -755,18 +752,14 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Install OpenAI + Anthropic Python SDKs run: pip install 'openai>=1.50' 'anthropic>=0.40' @@ -1040,6 +1033,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: json-images-log diff --git a/.github/workflows/studio-mac-install-matrix.yml b/.github/workflows/studio-mac-install-matrix.yml new file mode 100644 index 0000000000..da944d4b5c --- /dev/null +++ b/.github/workflows/studio-mac-install-matrix.yml @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Proves Studio's llama.cpp install loads on every supported macOS. The heavy +# app smokes stay single-OS; this matrix covers the OS-version dimension cheaply +# (install.sh + binary-load assert). Regression guard for the macOS-version +# selection in studio/install_llama_prebuilt.py. + +name: Mac Studio Install Matrix CI + +on: + pull_request: + paths: + - 'studio/install_llama_prebuilt.py' + - 'studio/setup.sh' + - 'install.sh' + - '.github/scripts/assert-llama-loads.sh' + - '.github/workflows/studio-mac-install-matrix.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + install-load: + name: Install + load (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + timeout-minutes: 25 + continue-on-error: ${{ matrix.experimental }} + strategy: + fail-fast: false + matrix: + include: + - os: macos-14 # Apple Silicon, macOS 14 Sonoma + experimental: false + - os: macos-15 # Apple Silicon, macOS 15 Sequoia + experimental: false + - os: macos-26 # Apple Silicon, macOS 26 Tahoe + experimental: false + - os: macos-15-intel # Intel x86_64, macOS 15 (informational) + experimental: true + - os: macos-26-intel # Intel x86_64, macOS 26 (last Intel macOS) + experimental: true + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: Upload install log + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: mac-install-matrix-${{ matrix.os }}-log + path: logs/install.log + retention-days: 7 diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index b353f0ec83..4f9f94b534 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -62,7 +62,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -73,29 +73,26 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Install Playwright + Chromium # No --with-deps on Mac: that flag installs Linux apt packages. diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml index b65439f174..f554a16415 100644 --- a/.github/workflows/studio-mac-update-smoke.yml +++ b/.github/workflows/studio-mac-update-smoke.yml @@ -62,30 +62,19 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - # Mac install must take the prebuilt path. Source-build - # fallback here is an Unsloth bug. - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi - if ! grep -qE "prebuilt installed and validated|prebuilt up to date and validated|bin-macos-arm64" logs/install.log; then - echo "::error::no Mac prebuilt llama.cpp marker in install.log." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi - echo "install.sh installed the Mac prebuilt llama.cpp" + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: First update should be a no-op (prebuilt already validated) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -104,6 +93,7 @@ jobs: - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 455fe4b7e1..de106e201f 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -76,7 +76,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -87,17 +87,19 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml index 057aeacbd4..307bb51972 100644 --- a/.github/workflows/studio-update-smoke.yml +++ b/.github/workflows/studio-update-smoke.yml @@ -71,6 +71,7 @@ jobs: # prebuilt path falls back to source build. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -85,6 +86,7 @@ jobs: # idempotency regressed. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -107,6 +109,7 @@ jobs: # the first one. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml index eee61516c7..78efe918ac 100644 --- a/.github/workflows/studio-windows-api-smoke.yml +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -69,7 +69,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -80,13 +80,14 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) shell: pwsh @@ -123,6 +124,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 6253b9d213..a772a6d102 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -13,7 +13,7 @@ # 2. Tool calling Tests # Qwen3.5-2B UD-Q4_K_XL (~890 MiB). # 3. JSON, images -# gemma-4-E2B-it UD-Q4_K_XL + mmproj-F16 (~3.4 GiB total). +# Qwen3-VL-2B-Instruct UD-IQ2_XXS + mmproj-F16 (~1.4 GiB total). # Within the 14 GB windows-latest SSD budget. name: Windows Studio GGUF CI @@ -65,6 +65,18 @@ jobs: with: persist-credentials: false + # Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit + # test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke. + - name: setup.ps1 unit test (Resolve-CudaToolkit) + shell: pwsh + run: | + $errs = $null + [void][System.Management.Automation.Language.Parser]::ParseFile( + (Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs) + if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } + Write-Host "setup.ps1 parsed with no errors" + pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1 + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' @@ -89,7 +101,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -102,6 +114,7 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME cache for ${{ env.GGUF_REPO }} # Only write a fresh cache entry when we actually rebuilt the @@ -111,7 +124,7 @@ jobs: uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) shell: pwsh @@ -148,6 +161,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -360,6 +374,9 @@ jobs: - name: Collect llama-server logs if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true shell: bash # Copy llama-server's own stdout/stderr (teed by Studio under # ~/.unsloth/studio/logs/llama-server/) into the workspace so @@ -373,6 +390,8 @@ jobs: - name: Upload logs if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: windows-openai-anthropic-log @@ -487,6 +506,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -788,6 +808,9 @@ jobs: - name: Collect llama-server logs if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true shell: bash # Copy llama-server's own stdout/stderr (teed by Studio under # ~/.unsloth/studio/logs/llama-server/) into the workspace so @@ -801,6 +824,8 @@ jobs: - name: Upload logs if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: windows-tool-calling-log @@ -821,9 +846,9 @@ jobs: run: shell: bash env: - GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF - GGUF_VARIANT: UD-Q4_K_XL - GGUF_FILE: gemma-4-E2B-it-UD-Q4_K_XL.gguf + GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF + GGUF_VARIANT: UD-IQ2_XXS + GGUF_FILE: Qwen3-VL-2B-Instruct-UD-IQ2_XXS.gguf MMPROJ_FILE: mmproj-F16.gguf STUDIO_PORT: '18899' HF_HOME: ${{ github.workspace }}/hf-cache @@ -857,7 +882,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - name: Prime HF_HOME with the GGUF + mmproj id: prime-hf @@ -869,13 +894,14 @@ jobs: mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME cache for ${{ env.GGUF_REPO }} (model + mmproj) if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) shell: pwsh @@ -912,6 +938,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -1101,7 +1128,7 @@ jobs: ) data_uri = f"data:image/png;base64,{PNG_64X64_RED_B64}" - # On Windows + the gemma-4-E2B mmproj, llama.cpp's vision + # On Windows + the Qwen3-VL mmproj, llama.cpp's vision # path runs on CPU (no Metal involvement). The wrapper is # kept for resilience but the vision path is expected to # work on Windows; an exception here is a real regression. @@ -1186,6 +1213,9 @@ jobs: - name: Collect llama-server logs if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true shell: bash # Copy llama-server's own stdout/stderr (teed by Studio under # ~/.unsloth/studio/logs/llama-server/) into the workspace so @@ -1199,6 +1229,8 @@ jobs: - name: Upload logs if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: windows-json-images-log diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index e3c8642122..40d8e530cd 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -85,7 +85,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -96,13 +96,14 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) shell: pwsh @@ -143,6 +144,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 redirects ALL PowerShell streams (stdout, stderr, diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index aa3e35f052..4a4806cfb1 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -133,6 +133,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -179,6 +180,7 @@ jobs: - name: First update should be a no-op (prebuilt already validated) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -197,6 +199,7 @@ jobs: - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1919fac9c5..cffbf73cd5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.14 + rev: v0.15.16 hooks: - id: ruff args: @@ -14,5 +14,20 @@ repos: entry: scripts/run_ruff_format.py language: python types: [python] + # Mirror ruff's [tool.ruff] extend-exclude so this hook does not + # half-process files ruff itself skips (which produced churn). + exclude: '(chat_templates|ollama_template_mappers|_auto_install|mapper)\.py$' additional_dependencies: - ruff==0.6.9 + # Re-pins allowScripts entries after dependency bumps. pre-commit.ci + # pushes the fix to PR branches, Dependabot's included, so stale pins + # heal without a human in the loop. + - id: sync-allow-scripts-pins + name: Sync allowScripts pins with the frontend lockfile + # `python + +""" + + +@studio_router.get("/artifact-preview-frame", include_in_schema = False) +async def artifact_preview_frame( + request: Request, + allow_network: bool = False, + token: Optional[str] = None, +): + """Serve the opaque sandbox shell used for client-side HTML artifacts.""" + + if allow_network: + auth_header = request.headers.get("authorization") + if auth_header and auth_header.lower().startswith("bearer "): + jwt_token = auth_header[7:] + elif token: + jwt_token = token + else: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Missing authentication token", + ) + from fastapi.security import HTTPAuthorizationCredentials + + creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = jwt_token) + await get_current_subject(creds) + + csp = ( + _ARTIFACT_PREVIEW_FRAME_NETWORK_CSP if allow_network else _ARTIFACT_PREVIEW_FRAME_STRICT_CSP + ) + return Response( + content = _ARTIFACT_PREVIEW_FRAME_HTML, + media_type = "text/html; charset=utf-8", + headers = { + "Cache-Control": "no-store", + "Content-Security-Policy": csp, + "Referrer-Policy": "no-referrer", + "X-Content-Type-Options": "nosniff", + }, + ) + + def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: - """Classify reasoning/tool capabilities via the GGUF classifier so - flags match across backends. gpt-oss is overridden because Harmony - routes reasoning and tools through tokenizer channels, not template - markup.""" + """Classify reasoning/tool capabilities via the GGUF classifier so flags + match across backends. gpt-oss is overridden: Harmony routes reasoning and + tools through tokenizer channels, not template markup.""" model_id = getattr(backend, "active_model_name", None) flags = ( detect_reasoning_flags( @@ -260,11 +867,11 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: "supports_tools": False, } ) - # Our safetensors loop only parses {json} - # and .... Llama uses <|python_tag|>, - # Mistral uses [TOOL_CALLS]; advertising tools for those would - # enable a pill the parser cannot honour. GGUF is unaffected -- - # llama-server normalises every format into structured deltas. + # Our safetensors loop only parses {json} and + # .... Llama uses <|python_tag|>, Mistral uses + # [TOOL_CALLS]; advertising tools for those enables a pill the parser + # can't honour. GGUF is unaffected -- llama-server normalises every + # format into structured deltas. if ( flags.get("supports_tools") and chat_template @@ -278,7 +885,7 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: ) flags["supports_tools"] = False - # gpt-oss: keep reasoning on, drop tools (Harmony channel, not + # gpt-oss: keep reasoning on, drop tools (Harmony channel, not the # XML this loop parses). try: if hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model(): @@ -294,7 +901,7 @@ def _effective_enable_tools(payload) -> Optional[bool]: """Resolve `payload.enable_tools` against the process-level tool policy. Returns the policy value when set (CLI hard-override from `unsloth run`), - otherwise the per-request value. + else the per-request value. """ from state.tool_policy import get_tool_policy @@ -302,23 +909,21 @@ def _effective_enable_tools(payload) -> Optional[bool]: return policy if policy is not None else payload.enable_tools -# Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts -# so is_disconnected() never fires. POST /inference/cancel looks up -# in-flight cancel_events here by cancel_id (per-run) or session_id / -# completion_id (fallbacks). +# Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts so +# is_disconnected() never fires. POST /inference/cancel looks up in-flight +# cancel_events here by cancel_id (per-run) or session_id / completion_id +# (fallbacks). _CANCEL_REGISTRY: dict[str, set[threading.Event]] = {} _CANCEL_LOCK = threading.Lock() -# Cancel POSTs that arrive before registration are stashed; the next -# matching __enter__ replays set() within the TTL. +# Cancel POSTs arriving before registration are stashed; the next matching +# __enter__ replays set() within the TTL. _PENDING_CANCELS: dict[str, float] = {} _PENDING_CANCEL_TTL_S = 30.0 def _prune_pending(now: float) -> None: - for k in [ - k for k, ts in _PENDING_CANCELS.items() if now - ts > _PENDING_CANCEL_TTL_S - ]: + for k in [k for k, ts in _PENDING_CANCELS.items() if now - ts > _PENDING_CANCEL_TTL_S]: _PENDING_CANCELS.pop(k, None) @@ -330,8 +935,8 @@ class _TrackedCancel: self.keys = tuple(k for k in keys if k) def __enter__(self): - # Register + consume-pending must be one critical section to close - # the TOCTOU race against a concurrent cancel POST. + # Register + consume-pending in one critical section to close the + # TOCTOU race against a concurrent cancel POST. should_cancel = False with _CANCEL_LOCK: for k in self.keys: @@ -359,9 +964,9 @@ class _TrackedCancel: def _cancel_by_keys(keys) -> int: """Set cancel_event for matching registry entries; no stash. - session_id/completion_id are shared across runs on the same thread, - so stashing them would ghost-cancel the user's next request. Only - cancel_id is per-run unique (see _cancel_by_cancel_id_or_stash).""" + session_id/completion_id are shared across runs on the same thread, so + stashing them would ghost-cancel the user's next request. Only cancel_id + is per-run unique (see _cancel_by_cancel_id_or_stash).""" if not keys: return 0 events: set[threading.Event] = set() @@ -396,17 +1001,15 @@ def _cancel_by_cancel_id_or_stash(cancel_id: str) -> int: async def _await_cancel_then_close(cancel_event, resp) -> None: """Watch a threading.Event from asyncio and close ``resp`` when it fires. - Used by the passthrough streamers so a /cancel POST can interrupt - while the async iterator is blocked waiting for llama-server prefill. - Without this watcher the in-loop ``cancel_event.is_set()`` check is - unreachable until the first SSE chunk arrives, which is exactly the - proxy/Colab scenario the cancel POST exists to handle. + Used by passthrough streamers so a /cancel POST can interrupt while the + async iterator is blocked on llama-server prefill. Without it the in-loop + ``cancel_event.is_set()`` check is unreachable until the first SSE chunk + arrives -- exactly the proxy/Colab case the cancel POST exists for. - Polls a threading.Event because the cancel registry is keyed by - threading.Event so the synchronous /cancel handler can call .set(). - 50ms cadence adds at most that much latency to a prefill cancel; the - common-case streaming cancel path still observes the event in the - iterator's first iteration after the next chunk. + Polls a threading.Event since the cancel registry is keyed by + threading.Event (so the sync /cancel handler can call .set()). The 50ms + cadence adds at most that latency to a prefill cancel; the common + streaming-cancel path still sees the event on the iterator's next chunk. """ try: while not cancel_event.is_set(): @@ -419,13 +1022,64 @@ async def _await_cancel_then_close(cancel_event, resp) -> None: return -# Appended to tool-use nudge to discourage plan-without-action -_TOOL_ACTION_NUDGE = ( - " IMPORTANT: Always call tools directly -- never write code yourself." - " Never describe what you plan to do -- just call the tool immediately." - " For any code request, call the python tool. For any factual question, call web_search." - " Do NOT output code blocks -- use the python tool instead." +# Centralized local/server tool nudge. Keep render_html guidance gated to turns +# where the artifact tool is actually present in the tool schema; otherwise +# small local models can hallucinate a missing tool call instead of following +# the fenced-HTML fallback prompt. +_TOOL_BASE_NUDGE = ( + "Tools are available when they materially improve the answer. Use an enabled " + "tool for current facts, calculations, code execution, or artifacts when it " + "materially helps; otherwise answer normally and follow the user's requested " + "format." ) +_TOOL_WEB_COMPACT_TIP = "When using web_search, do not repeat the same search query." +_TOOL_WEB_EXPANDED_TIP = ( + "When using web_search and a result URL is relevant, fetch its full content " + "by calling web_search with the url parameter. Do not repeat the same search " + "query. If a search returns no useful results, try rephrasing or fetching a " + "result URL directly." +) +_TOOL_CODE_TIP = ( + "Use code execution for math, calculations, data processing, or to parse " + "and analyze information from tool results." +) +_TOOL_ARTIFACT_TIP = ( + "For HTML, CSS, or JavaScript artifact requests, call render_html once when " + "it is available with one complete self-contained HTML document in the code " + "argument. After render_html succeeds, do not call it again in the same " + "response unless the user asks for changes. Future user requests for new " + "artifacts may call render_html once." +) + + +def _build_tool_action_nudge(*, tools: list[dict], model_name: str) -> str: + tool_names = { + (tool.get("function") or {}).get("name") + for tool in tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + } + has_web = "web_search" in tool_names + has_code = "python" in tool_names or "terminal" in tool_names + has_artifact = "render_html" in tool_names + if not (has_web or has_code or has_artifact): + return "" + + model_size_b = _extract_model_size_b(model_name) + compact_web_tip = model_size_b is not None and model_size_b < 9 + tool_tip_parts: list[str] = [] + if has_web: + tool_tip_parts.append(_TOOL_WEB_COMPACT_TIP if compact_web_tip else _TOOL_WEB_EXPANDED_TIP) + if has_code: + tool_tip_parts.append(_TOOL_CODE_TIP) + if has_artifact: + tool_tip_parts.append(_TOOL_ARTIFACT_TIP) + return ( + f"The current date is {_date.today().isoformat()}. " + + _TOOL_BASE_NUDGE + + " " + + " ".join(tool_tip_parts) + ) + # Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py # split across the visible/DRAIN boundary. Four leak shapes: @@ -435,47 +1089,61 @@ _TOOL_ACTION_NUDGE = ( # 4. tail-only `` (outer close truncated by EOS); anchored to # `\Z` so mid-text `` in user code samples survives. _TOOL_XML_RE = _re.compile( - r"<(?:tool_call|function=\w+)>.*?(?:|\Z)" + # Hyphen in the name char-class matches MCP tool names with dashes + # (mcp__srv__list-issues) that would otherwise leak past this strip. + r"<(?:tool_call|function=[\w-]+)>.*?(?:|\Z)" r"|" r"|\s*\Z", _re.DOTALL, ) + + +def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str: + """Apply route-level XML leak cleanup only when Auto-Heal is enabled.""" + if not auto_heal_tool_calls: + return text + return _TOOL_XML_RE.sub("", text) + + logger = get_logger(__name__) -def _validate_native_mmproj_companion( - mmproj_path: str | None, gguf_path: str | None +def _validate_native_gguf_companion( + companion_path: str | None, gguf_path: str | None, label: str ) -> None: - if not mmproj_path or not gguf_path: + """Reject a companion GGUF (mmproj / MTP drafter) that a native-lease load + would otherwise hand to llama-server: must be a regular file (no symlink + escaping the leased directory) living next to the selected GGUF.""" + if not companion_path or not gguf_path: return import stat as _stat_module - mm = Path(mmproj_path) + companion = Path(companion_path) gguf = Path(gguf_path) try: - mm_lstat = os.lstat(mm) + companion_lstat = os.lstat(companion) except OSError as exc: raise HTTPException( status_code = 400, - detail = "Native vision companion is no longer accessible.", + detail = f"Native {label} is no longer accessible.", ) from exc - if _stat_module.S_ISLNK(mm_lstat.st_mode) or not _stat_module.S_ISREG( - mm_lstat.st_mode + if _stat_module.S_ISLNK(companion_lstat.st_mode) or not _stat_module.S_ISREG( + companion_lstat.st_mode ): raise HTTPException( status_code = 400, - detail = "Native vision companion must be a regular file.", + detail = f"Native {label} must be a regular file.", ) try: - if mm.resolve(strict = True).parent != gguf.resolve(strict = True).parent: + if companion.resolve(strict = True).parent != gguf.resolve(strict = True).parent: raise HTTPException( status_code = 400, - detail = "Native vision companion must live next to the selected GGUF.", + detail = f"Native {label} must live next to the selected GGUF.", ) except OSError as exc: raise HTTPException( status_code = 400, - detail = "Native vision companion is no longer accessible.", + detail = f"Native {label} is no longer accessible.", ) from exc @@ -489,42 +1157,36 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: 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. +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_mode = "off" - else: - req_mode = _canonicalize_spec_mode(request.speculative_type) or "auto" + # Spec decoding works on vision models too (MTP is mmproj-compatible, + # llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare + # the real requested mode -- coercing vision to ``off`` here used to + # swallow every spec-mode change on a vision model as already_loaded. + req_mode = _canonicalize_spec_mode(request.speculative_type) or "auto" backend_mode = llama_backend.requested_spec_mode or "auto" if req_mode != backend_mode: return False - # spec_draft_n_max only matters when an MTP variant is engaged; None - # means "platform default" and matches whatever the backend chose. + # spec_draft_n_max only matters with an MTP variant; None means "platform + # default" and matches whatever the backend chose. if backend_mode in ("mtp", "mtp+ngram") and request.spec_draft_n_max is not None: if int(request.spec_draft_n_max) != (llama_backend.spec_draft_n_max or 0): return False - if (request.chat_template_override or None) != ( - llama_backend.chat_template_override or None - ): + 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. + # llama_extra_args=None means "inherit"; only an explicit differing list + # forces a reload. On the inherit path, refuse to match if stored extras + # contain any shadow flag, so the reload path strips them rather than + # 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: @@ -532,13 +1194,38 @@ def _request_matches_loaded_settings( else: if list(request.llama_extra_args) != backend_extra: return False + # A separate drafter (Gemma's root mtp-*.gguf) appearing or disappearing + # next to the loaded weights changes the launch command (--model-draft), + # so a duplicate /load must reload rather than dedupe. Always compare the + # detected vs stored drafter when the mode can use one and the user does + # not own --spec-type: the resolved-path compare is cheap and handles all + # four cases (both None -> match; one None -> reload; equal -> match; + # different -> reload), including a drafter deleted out from under a + # running server. Runs last: it stats the filesystem, so every pure-memory + # comparison above short-circuits first. Resolve both sides since the + # stored launch path may be a snapshot symlink while detect_mtp_file + # returns the resolved blob. + if req_mode in ("auto", "mtp", "mtp+ngram") and llama_backend.gguf_path: + effective_extras = ( + request.llama_extra_args + if request.llama_extra_args is not None + else llama_backend.extra_args + ) + if not _extra_args_set_spec_type(effective_extras): + detected = detect_mtp_file(llama_backend.gguf_path) + stored = llama_backend.mtp_draft_path + try: + detected_resolved = Path(detected).resolve() if detected else None + stored_resolved = Path(stored).resolve() if stored else None + except OSError: + return False + if detected_resolved != stored_resolved: + return False return True def _resolve_model_identifier_for_request( - request: LoadRequest | ValidateModelRequest, - *, - operation: str, + request: LoadRequest | ValidateModelRequest, *, operation: str ) -> tuple[str, str, bool]: if not request.native_path_lease: return request.model_path, request.model_path, False @@ -551,10 +1238,14 @@ def _resolve_model_identifier_for_request( allowed_suffixes = (".gguf",), ) except NativePathLeaseError as exc: - raise HTTPException(status_code = 400, detail = str(exc)) from exc - display_label = ( - grant.display_label or Path(request.model_path).name or "Native model" - ) + # Curated, client-correctable lease error (expired / wrong type / re-select); + # keep the actionable message, just redact paths. + logger.warning("inference.native_path_lease_failed: %s", exc) + raise HTTPException( + status_code = 400, + detail = redact_native_paths(str(exc)), + ) from exc + display_label = grant.display_label or Path(request.model_path).name or "Native model" return str(grant.canonical_path), display_label, True @@ -575,23 +1266,28 @@ async def load_model( """ Load a model for inference. - The model_path should be a clean identifier from GET /models/list. - Returns inference configuration parameters (temperature, top_p, top_k, min_p) - from the model's YAML config, falling back to default.yaml for missing values. + model_path is a clean identifier from GET /models/list. Returns inference + config (temperature, top_p, top_k, min_p) from the model's YAML, falling + back to default.yaml for missing values. - GGUF models are loaded via llama-server (llama.cpp) instead of Unsloth. + GGUF models load via llama-server (llama.cpp) instead of Unsloth. """ native_grant_backed = False model_log_label = request.model_path try: - # Validate user-supplied llama-server pass-through args up front - # so a managed-flag collision returns 400 before any model work. + # Validate user pass-through args up front so a managed-flag collision + # returns 400 before any model work. try: 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 []". + # Keep the curated validation message (names the flag); just strip paths. + logger.warning("inference.validate_extra_args_failed: %s", exc) + raise HTTPException( + status_code = 400, + detail = redact_native_paths(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 ) @@ -599,8 +1295,8 @@ async def load_model( model_identifier, model_log_label, native_grant_backed = ( _resolve_model_identifier_for_request(request, operation = "load-model") ) - # Version switching is handled automatically by the subprocess-based - # inference backend — no need for ensure_transformers_version() here. + # Version switching is handled by the subprocess-based inference + # backend -- no ensure_transformers_version() needed here. # ── Already-loaded check: skip reload if the exact model is active ── backend = get_inference_backend() @@ -618,7 +1314,7 @@ async def load_model( and gguf_variant_matches and llama_backend.model_identifier and llama_backend.model_identifier.lower() == model_identifier.lower() - # Match runtime settings too so Apply isn't dropped (#5401). + # Match runtime settings so Apply isn't dropped (#5401). and _request_matches_loaded_settings(request, llama_backend) # Skip if a prior audio probe failed -- let load_model retry. and getattr(llama_backend, "_audio_probed", True) @@ -630,9 +1326,7 @@ async def load_model( inference_config = load_inference_config(llama_backend.model_identifier) _gguf_audio = ( - llama_backend._audio_type - if hasattr(llama_backend, "_audio_type") - else None + llama_backend._audio_type if hasattr(llama_backend, "_audio_type") else None ) _gguf_is_audio = getattr(llama_backend, "_is_audio", False) return LoadResponse( @@ -648,7 +1342,7 @@ async def load_model( is_gguf = True, is_audio = _gguf_is_audio, audio_type = _gguf_audio, - has_audio_input = False, + has_audio_input = getattr(llama_backend, "_has_audio_input", False), inference = inference_config, requires_trust_remote_code = bool( inference_config.get("trust_remote_code", False) @@ -670,9 +1364,7 @@ async def load_model( backend.active_model_name and backend.active_model_name.lower() == model_identifier.lower() ): - logger.info( - f"Model already loaded (Unsloth): {model_log_label}, skipping reload" - ) + logger.info(f"Model already loaded (Unsloth): {model_log_label}, skipping reload") inference_config = load_inference_config(backend.active_model_name) _model_info = backend.models.get(backend.active_model_name, {}) _chat_template = None @@ -689,9 +1381,7 @@ async def load_model( _sf_reasoning_style = _sf_flags["reasoning_style"] return LoadResponse( status = "already_loaded", - model = model_log_label - if native_grant_backed - else backend.active_model_name, + model = model_log_label if native_grant_backed else backend.active_model_name, display_name = model_log_label if native_grant_backed else backend.active_model_name, @@ -710,12 +1400,13 @@ async def load_model( reasoning_always_on = _sf_flags["reasoning_always_on"], supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], supports_tools = _sf_flags["supports_tools"], + context_length = _positive_int_or_none(_model_info.get("context_length")), chat_template = _chat_template, ) # is_lora auto-detected from adapter_config.json on disk/HF. - # DNS-probe wrap so offline loads skip 30-60s of soft-failed - # network checks before the worker starts. + # DNS-probe wrap so offline loads skip 30-60s of soft-failed network + # checks before the worker starts. with _hf_offline_if_dns_dead(): config = ModelConfig.from_identifier( model_id = model_identifier, @@ -743,28 +1434,26 @@ async def load_model( llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() - # Unload any active Unsloth model first to free VRAM + # Unload any active Unsloth model to free VRAM if unsloth_backend.active_model_name: logger.info( f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" ) 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. + # Inherit llama_extra_args from the previous load when the request + # omits the field (the chat-settings Apply path doesn't round-trip + # them; explicit [] still clears). 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. + # 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 HF / local branches + # below), so both sides key off the same string. resolved_variant = config.gguf_variant same_source = bool( source @@ -774,35 +1463,32 @@ async def load_model( ) if not same_source: logger.info( - "Not inheriting llama_extra_args: stored args came " - "from %s, loading %s", + "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. + # 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. + # Strip only the groups whose first-class field was 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 - or "spec_draft_n_max" in fields_set + "speculative_type" in fields_set or "spec_draft_n_max" 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. + # Shouldn't 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", @@ -817,10 +1503,9 @@ async def load_model( 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) - # GGUF download + llama-server startup. + # Route to HF 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) GGUF download + llama-server start. _n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1) if config.gguf_hf_repo: @@ -842,17 +1527,29 @@ async def load_model( ) else: # Local mode: llama-server loads via -m - if native_grant_backed and config.gguf_mmproj_file: - _validate_native_mmproj_companion( - config.gguf_mmproj_file, config.gguf_file - ) + if native_grant_backed: + if config.gguf_mmproj_file: + _validate_native_gguf_companion( + config.gguf_mmproj_file, config.gguf_file, "vision companion" + ) + if config.gguf_mtp_file: + # The drafter is optional (unlike mmproj for a vision + # model): drop it rather than fail the load. + try: + _validate_native_gguf_companion( + config.gguf_mtp_file, config.gguf_file, "MTP drafter" + ) + except HTTPException as exc: + logger.warning("Dropping MTP drafter for native load: %s", exc.detail) + config.gguf_mtp_file = None success = await asyncio.to_thread( 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). + mtp_draft_path = config.gguf_mtp_file, + # Pass the resolved variant so _extra_args_source keys 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, @@ -878,9 +1575,7 @@ async def load_model( # Audio detection moved into load_model under _serial_load_lock (#5642). _gguf_audio = llama_backend._audio_type _gguf_is_audio = llama_backend._is_audio - llama_backend._native_display_label = ( - model_log_label if native_grant_backed else None - ) + llama_backend._native_display_label = model_log_label if native_grant_backed else None llama_backend._native_grant_backed = bool(native_grant_backed) if _gguf_is_audio: logger.info(f"GGUF model detected as audio: audio_type={_gguf_audio}") @@ -890,19 +1585,15 @@ async def load_model( return LoadResponse( status = "loaded", model = model_log_label if native_grant_backed else config.identifier, - display_name = model_log_label - if native_grant_backed - else config.display_name, + display_name = model_log_label if native_grant_backed else config.display_name, is_vision = llama_backend.is_vision, is_lora = False, is_gguf = True, is_audio = _gguf_is_audio, audio_type = _gguf_audio, - has_audio_input = False, + has_audio_input = llama_backend._has_audio_input, inference = inference_config, - requires_trust_remote_code = bool( - inference_config.get("trust_remote_code", False) - ), + requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)), context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, @@ -929,12 +1620,9 @@ async def load_model( # Shut down any export subprocess to free VRAM try: from core.export import get_export_backend - exp_backend = get_export_backend() if exp_backend.current_checkpoint: - logger.info( - "Shutting down export subprocess to free GPU memory for inference" - ) + logger.info("Shutting down export subprocess to free GPU memory for inference") exp_backend._shutdown_subprocess() exp_backend.current_checkpoint = None exp_backend.is_vision = False @@ -942,9 +1630,9 @@ async def load_model( except Exception as e: logger.warning("Could not shut down export subprocess: %s", e) - # Auto-detect quantization for LoRA adapters from adapter_config.json - # The training pipeline patches this file with "unsloth_training_method" - # which is 'qlora' or 'lora'. Only LoRA (16-bit) needs load_in_4bit=False. + # Auto-detect quantization for LoRA adapters from adapter_config.json. + # The training pipeline writes "unsloth_training_method" ('qlora' or + # 'lora'); only LoRA (16-bit) needs load_in_4bit=False. load_in_4bit = request.load_in_4bit if config.is_lora and config.path: import json @@ -973,7 +1661,7 @@ async def load_model( f"Training method: {training_method}, load_in_4bit={load_in_4bit}" ) else: - # No unsloth_training_method — fallback to base model name + # No unsloth_training_method -- fall back to base model name if ( config.base_model and "-bnb-4bit" not in config.base_model.lower() @@ -988,8 +1676,8 @@ async def load_model( except Exception as e: logger.warning(f"Could not read adapter_config.json: {e}") - # Load the model in a thread so the event loop stays free - # for download progress polling and other requests. + # Load in a thread so the event loop stays free for download progress + # polling and other requests. success = await asyncio.to_thread( backend.load_model, config = config, @@ -1001,12 +1689,10 @@ async def load_model( ) if not success: - # Check if YAML says this model needs trust_remote_code + # Check if YAML says this model needs trust_remote_code. if not request.trust_remote_code: model_defaults = load_model_defaults(config.identifier) - yaml_trust = model_defaults.get("inference", {}).get( - "trust_remote_code", False - ) + yaml_trust = model_defaults.get("inference", {}).get("trust_remote_code", False) if yaml_trust: raise HTTPException( status_code = 400, @@ -1042,9 +1728,7 @@ async def load_model( return LoadResponse( status = "loaded", model = model_log_label if native_grant_backed else config.identifier, - display_name = model_log_label - if native_grant_backed - else config.display_name, + display_name = model_log_label if native_grant_backed else config.display_name, is_vision = config.is_vision, is_lora = config.is_lora, is_gguf = False, @@ -1052,14 +1736,13 @@ async def load_model( audio_type = config.audio_type, has_audio_input = config.has_audio_input, inference = inference_config, - requires_trust_remote_code = bool( - inference_config.get("trust_remote_code", False) - ), + requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)), supports_reasoning = _sf_flags["supports_reasoning"], reasoning_style = _sf_flags["reasoning_style"], reasoning_always_on = _sf_flags["reasoning_always_on"], supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], supports_tools = _sf_flags["supports_tools"], + context_length = _positive_int_or_none(_model_info.get("context_length")), chat_template = _chat_template, ) @@ -1075,9 +1758,10 @@ async def load_model( ) raise HTTPException(status_code = 400, detail = redacted_msg) logger.warning("Rejected inference GPU selection: %s", e) - raise HTTPException(status_code = 400, detail = str(e)) + # User-facing validation (e.g. "Invalid gpu_ids [99]"): redact paths, keep detail. + raise HTTPException(status_code = 400, detail = redact_native_paths(str(e))) except Exception as e: - # Surface a friendlier message for models that Unsloth cannot load + # Friendlier message for models Unsloth cannot load. not_supported_hints = [ "No config file found", "not yet supported", @@ -1099,7 +1783,7 @@ async def load_model( detail = f"Failed to load native model {model_log_label}: {msg}", ) logger.error(f"Error loading model: {e}", exc_info = True) - msg = str(e) + msg = redact_native_paths(str(e)) if any(h.lower() in msg.lower() for h in not_supported_hints): msg = f"This model is not supported yet. Try a different model. (Original error: {msg})" raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") @@ -1107,14 +1791,13 @@ async def load_model( @router.post("/validate", response_model = ValidateModelResponse) async def validate_model( - request: ValidateModelRequest, - current_subject: str = Depends(get_current_subject), + request: ValidateModelRequest, current_subject: str = Depends(get_current_subject) ): """ Lightweight validation endpoint for model identifiers. - This checks that ModelConfig.from_identifier() can resolve the given - model_path, but it does NOT actually load model weights into GPU memory. + Checks that ModelConfig.from_identifier() can resolve model_path, but does + NOT load model weights into GPU memory. """ native_grant_backed = False model_log_label = request.model_path @@ -1178,27 +1861,22 @@ async def validate_model( ) raise HTTPException( status_code = 400, - detail = f"Invalid model: {str(e)}", + detail = "Invalid model", ) @router.post("/unload", response_model = UnloadResponse) -async def unload_model( - request: UnloadRequest, - current_subject: str = Depends(get_current_subject), -): +async def unload_model(request: UnloadRequest, current_subject: str = Depends(get_current_subject)): """ Unload a model from memory. Routes to the correct backend (llama-server for GGUF, Unsloth otherwise). """ try: - # Check if the GGUF backend has this model loaded or is loading it + # Check if the GGUF backend has this model loaded or is loading it. llama_backend = get_llama_cpp_backend() if llama_backend.is_active and ( llama_backend.model_identifier == request.model_path - or is_registered_native_path_label( - llama_backend.model_identifier, request.model_path - ) + or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) or not llama_backend.is_loaded ): llama_backend.unload_model() @@ -1213,14 +1891,11 @@ async def unload_model( except Exception as e: logger.error(f"Error unloading model: {e}", exc_info = True) - raise HTTPException(status_code = 500, detail = f"Failed to unload model: {str(e)}") + raise HTTPException(status_code = 500, detail = "Failed to unload model") @studio_router.post("/cancel") -async def cancel_inference( - request: Request, - current_subject: str = Depends(get_current_subject), -): +async def cancel_inference(request: Request, current_subject: str = Depends(get_current_subject)): """Cancel in-flight inference requests. Body (JSON, at least one key required): @@ -1228,8 +1903,8 @@ async def cancel_inference( session_id - fallback when cancel_id is absent. completion_id - fallback when cancel_id is absent. - A cancel_id arriving before its stream registers is stashed briefly - and replayed on registration. Returns {"cancelled": N}. + A cancel_id arriving before its stream registers is stashed briefly and + replayed on registration. Returns {"cancelled": N}. """ try: body = await request.json() @@ -1244,8 +1919,8 @@ async def cancel_inference( return {"cancelled": _cancel_by_cancel_id_or_stash(cancel_id)} keys = [] - # `message_id` is the Anthropic passthrough's per-run identifier -- - # included so /v1/messages clients can cancel by their native id. + # `message_id` is the Anthropic passthrough's per-run identifier, so + # /v1/messages clients can cancel by their native id. for k in ("completion_id", "session_id", "message_id"): v = body.get(k) if isinstance(v, str) and v: @@ -1260,13 +1935,12 @@ async def cancel_inference( @router.post("/generate/stream") async def generate_stream( - request: GenerateRequest, - current_subject: str = Depends(get_current_subject), + request: GenerateRequest, current_subject: str = Depends(get_current_subject) ): """ Generate a chat response with Server-Sent Events (SSE) streaming. - For vision models, provide image_base64 with the base64-encoded image. + For vision models, provide image_base64 (base64-encoded image). """ backend = get_inference_backend() @@ -1275,7 +1949,7 @@ async def generate_stream( status_code = 400, detail = "No model loaded. Call POST /inference/load first." ) - # Decode image if provided (for vision models) + # Decode image if provided (vision models) image = None if request.image_base64: try: @@ -1283,7 +1957,7 @@ async def generate_stream( from PIL import Image from io import BytesIO - # Check if current model supports vision + # Check current model supports vision model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_vision"): raise HTTPException( @@ -1298,8 +1972,12 @@ async def generate_stream( except HTTPException: raise except Exception as e: - raise HTTPException( - status_code = 400, detail = f"Failed to decode image: {str(e)}" + raise log_and_http_error( + e, + 400, + "Failed to decode image", + event = "inference.decode_image_failed", + log = logger, ) async def stream(): @@ -1333,17 +2011,15 @@ async def generate_stream( @router.get("/status", response_model = InferenceStatusResponse) -async def get_status( - current_subject: str = Depends(get_current_subject), -): +async def get_status(current_subject: str = Depends(get_current_subject)): """ Get current inference backend status. - Reports whichever backend (Unsloth or llama-server) is currently active. + Reports whichever backend (Unsloth or llama-server) is active. """ try: llama_backend = get_llama_cpp_backend() - # MTP probe + freshness check (both cached). Drive the UI banner. + # MTP probe + freshness check (both cached); drive the UI banner. try: _bin = type(llama_backend)._find_llama_server_binary() _caps = type(llama_backend).probe_server_capabilities(_bin) @@ -1353,7 +2029,6 @@ async def get_status( _supports_mtp = True # fail open try: from utils.llama_cpp_freshness import check_prebuilt_freshness - _freshness = check_prebuilt_freshness(_bin) except Exception: _freshness = {} @@ -1385,7 +2060,7 @@ async def get_status( gguf_variant = llama_backend.hf_variant, is_audio = getattr(llama_backend, "_is_audio", False), audio_type = _audio_type, - has_audio_input = False, + has_audio_input = getattr(llama_backend, "_has_audio_input", False), loading = [], loaded = [_display_model_id] if _display_model_id else [], inference = _inference_cfg, @@ -1406,6 +2081,7 @@ async def get_status( speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, llama_cpp_supports_mtp = _supports_mtp, + spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, llama_cpp_installed_tag = _installed_tag, llama_cpp_latest_tag = _latest_tag, @@ -1427,17 +2103,13 @@ async def get_status( has_audio_input = model_info.get("has_audio_input", False) chat_template_info = model_info.get("chat_template_info", {}) chat_template = ( - chat_template_info.get("template") - if isinstance(chat_template_info, dict) - else None + chat_template_info.get("template") if isinstance(chat_template_info, dict) else None ) # Non-GGUF: classify from the loaded template. _sf_flags = _detect_safetensors_features(backend, chat_template) inference_config = ( - load_inference_config(backend.active_model_name) - if backend.active_model_name - else None + load_inference_config(backend.active_model_name) if backend.active_model_name else None ) return InferenceStatusResponse( @@ -1459,6 +2131,7 @@ async def get_status( reasoning_always_on = _sf_flags["reasoning_always_on"], supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], supports_tools = _sf_flags["supports_tools"], + context_length = _positive_int_or_none(model_info.get("context_length")), chat_template = chat_template, llama_cpp_supports_mtp = _supports_mtp, llama_cpp_prebuilt_stale = _stale, @@ -1468,25 +2141,21 @@ async def get_status( except Exception as e: logger.error(f"Error getting status: {e}", exc_info = True) - raise HTTPException(status_code = 500, detail = f"Failed to get status: {str(e)}") + raise HTTPException(status_code = 500, detail = "Failed to get status") @router.get("/load-progress", response_model = LoadProgressResponse) -async def get_load_progress( - current_subject: str = Depends(get_current_subject), -): +async def get_load_progress(current_subject: str = Depends(get_current_subject)): """ Return the active GGUF load's mmap/upload progress. - During the warmup window after a GGUF download -- when llama-server - is paging ~tens-to-hundreds of GB of shards into the page cache - before pushing layers to VRAM -- ``/api/inference/status`` only - shows a generic spinner. This endpoint exposes sampled progress so - the UI can render a real bar plus rate/ETA during that window. + During the warmup window after a GGUF download -- when llama-server pages + ~tens-to-hundreds of GB of shards into the page cache before pushing layers + to VRAM -- ``/api/inference/status`` only shows a generic spinner. This + exposes sampled progress so the UI can render a real bar plus rate/ETA. - Returns an empty payload (``phase=null, bytes=0``) when no load is - in flight. The frontend should stop polling once ``phase`` becomes - ``ready``. + Returns an empty payload (``phase=null, bytes=0``) when no load is in + flight. The frontend should stop polling once ``phase`` becomes ``ready``. """ try: llama_backend = get_llama_cpp_backend() @@ -1512,7 +2181,7 @@ async def generate_audio( ): """ Generate audio (TTS) from the latest user message. - Returns a JSON response with base64-encoded WAV audio. + Returns JSON with base64-encoded WAV audio. Works with both GGUF (llama-server) and Unsloth/transformers backends. """ import base64 @@ -1521,9 +2190,7 @@ async def generate_audio( _, chat_messages, _ = _extract_content_parts(payload.messages) if not chat_messages: raise HTTPException(status_code = 400, detail = "No messages provided.") - last_user_msg = next( - (m for m in reversed(chat_messages) if m["role"] == "user"), None - ) + last_user_msg = next((m for m in reversed(chat_messages) if m["role"] == "user"), None) if not last_user_msg: raise HTTPException(status_code = 400, detail = "No user message found.") text = last_user_msg["content"] @@ -1539,7 +2206,7 @@ async def generate_audio( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, + max_new_tokens = _effective_max_tokens(payload) or 2048, repetition_penalty = payload.repetition_penalty, ) else: @@ -1548,9 +2215,7 @@ async def generate_audio( raise HTTPException(status_code = 400, detail = "No model loaded.") model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_audio"): - raise HTTPException( - status_code = 400, detail = "Active model is not an audio model." - ) + raise HTTPException(status_code = 400, detail = "Active model is not an audio model.") model_name = backend.active_model_name gen = lambda: backend.generate_audio_response( text = text, @@ -1558,18 +2223,16 @@ async def generate_audio( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, + max_new_tokens = _effective_max_tokens(payload) or 2048, repetition_penalty = payload.repetition_penalty, use_adapter = payload.use_adapter, ) try: - wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor( - None, gen - ) + wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor(None, gen) except Exception as e: logger.error(f"Audio generation error: {e}", exc_info = True) - raise HTTPException(status_code = 500, detail = str(e)) + raise HTTPException(status_code = 500, detail = safe_error_detail(e)) audio_b64 = base64.b64encode(wav_bytes).decode("ascii") return JSONResponse( @@ -1606,8 +2269,8 @@ def _decode_audio_base64(b64: str) -> np.ndarray: from utils.paths import ensure_dir, tmp_root raw = base64.b64decode(b64) - # torchaudio.load needs a file path or file-like object with format hint - # Write to a temp file so torchaudio can auto-detect the format + # torchaudio.load needs a path or file-like with a format hint; write a + # temp file so it can auto-detect the format. with tempfile.NamedTemporaryFile( suffix = ".audio", delete = False, @@ -1632,9 +2295,173 @@ def _decode_audio_base64(b64: str) -> np.ndarray: return waveform.squeeze(0).numpy() -def _extract_content_parts( - messages: list, -) -> tuple[str, list[dict], "Optional[str]"]: +# Reject oversized audio before decoding. base64 inflates raw bytes by ~4/3, so +# cap the encoded length to bound the upload. _MAX_AUDIO_SECONDS additionally +# bounds the *decoded* length, since a small compressed file (opus/flac/etc.) +# can expand to a far larger PCM array than the encoded-size cap implies. +_MAX_AUDIO_RAW_BYTES = 25 * 1024 * 1024 +_MAX_AUDIO_B64_CHARS = _MAX_AUDIO_RAW_BYTES * 4 // 3 +_MAX_AUDIO_SECONDS = 30 * 60 +_WAV_HEADER_BYTES = 44 +_MIN_TRANSCODE_AUDIO_SAMPLE_RATE = 8000 + + +def _sniff_audio_container(raw: bytes) -> Optional[str]: + """Return 'wav' or 'mp3' if the bytes are a container llama-server accepts + directly (so we can forward them untouched), else None (needs transcoding).""" + if len(raw) >= 12 and raw[:4] == b"RIFF" and raw[8:12] == b"WAVE": + return "wav" + # mp3: ID3 tag, or an MPEG audio frame sync (no other accepted format leads + # with 0xFF, so the simple sync check doesn't collide). + if raw[:3] == b"ID3" or (len(raw) >= 2 and raw[0] == 0xFF and (raw[1] & 0xE0) == 0xE0): + return "mp3" + return None + + +def _mono_f32_to_wav_bytes(arr: np.ndarray, sample_rate: int) -> bytes: + """Encode a mono float32 array as 16-bit PCM WAV bytes. + + Torch-free (numpy + stdlib only) so it works on no-torch GGUF-only installs; + the shared audio_codecs helper pulls in torch at import time. + """ + import io + import wave + + arr = np.nan_to_num(np.asarray(arr, dtype = np.float32).flatten(), posinf = 0.0, neginf = 0.0) + if arr.size == 0: + raise ValueError("decoded audio is empty") + peak = float(np.abs(arr).max()) + if peak > 1.0: + arr = arr / peak + pcm = (arr * 32767.0).astype(np.int16) + + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(int(sample_rate)) + wf.writeframes(pcm.tobytes()) + return buf.getvalue() + + +def _resample_mono_linear(arr: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray: + """Small numpy-only resampler for upload size limiting.""" + if source_rate <= 0 or target_rate <= 0 or source_rate == target_rate: + return arr + duration = len(arr) / float(source_rate) + target_len = max(1, int(round(duration * target_rate))) + if target_len == len(arr): + return arr + source_x = np.linspace(0.0, duration, num = len(arr), endpoint = False) + target_x = np.linspace(0.0, duration, num = target_len, endpoint = False) + return np.interp(target_x, source_x, arr).astype(np.float32) + + +def _fit_transcoded_audio_to_wav_cap(arr: np.ndarray, sample_rate: int) -> tuple[np.ndarray, int]: + """Downsample only when needed so transcoded WAV stays within the upload cap.""" + if sample_rate <= 0: + raise ValueError("decoded audio has an invalid sample rate") + wav_bytes = _WAV_HEADER_BYTES + len(arr) * 2 + if wav_bytes <= _MAX_AUDIO_RAW_BYTES: + return arr, sample_rate + + duration = len(arr) / float(sample_rate) + max_samples = max(1, (_MAX_AUDIO_RAW_BYTES - _WAV_HEADER_BYTES) // 2) + target_rate = int(max_samples // duration) + if target_rate < _MIN_TRANSCODE_AUDIO_SAMPLE_RATE: + raise ValueError("decoded audio exceeds the transcoded WAV size limit") + target_rate = min(sample_rate, target_rate) + fitted = _resample_mono_linear(arr, sample_rate, target_rate) + if _WAV_HEADER_BYTES + len(fitted) * 2 > _MAX_AUDIO_RAW_BYTES: + raise ValueError("decoded audio exceeds the transcoded WAV size limit") + return fitted, target_rate + + +def _decode_audio_mono(raw: bytes) -> tuple[np.ndarray, int]: + """Decode audio bytes to (mono float32 array, native sample_rate). + + soundfile (libsndfile) reads wav/mp3/ogg/flac straight from memory. librosa + (ffmpeg-backed) additionally covers m4a/webm but needs a real path and is + absent on no-torch GGUF-only installs. Both imports are inside the fallback + so a missing decoder degrades to the next one (and finally a clear error) + rather than crashing. + """ + import io + + try: + import soundfile as sf + arr, sr = sf.read(io.BytesIO(raw), dtype = "float32") + except Exception: + try: + import librosa + except ModuleNotFoundError as e: + raise RuntimeError( + "this audio format needs librosa, which is not installed in " + "GGUF-only environments; use wav, mp3, ogg or flac" + ) from e + import os + import tempfile + from utils.paths import ensure_dir, tmp_root + + with tempfile.NamedTemporaryFile( + suffix = ".audio", + delete = False, + dir = str(ensure_dir(tmp_root())), + ) as tmp: + tmp.write(raw) + tmp_path = tmp.name + try: + arr, sr = librosa.load(tmp_path, sr = None, mono = True) + finally: + os.unlink(tmp_path) + if arr.ndim > 1: + arr = arr.mean(axis = 1) + if sr > 0 and len(arr) > sr * _MAX_AUDIO_SECONDS: + raise ValueError(f"decoded audio exceeds the {_MAX_AUDIO_SECONDS // 60}-minute limit") + return arr, sr + + +def _prepare_audio_for_llama(b64: str) -> tuple[str, str]: + """Return (base64, format) ready for llama-server's input_audio part. + + llama-server's API only accepts wav/mp3, and decodes/resamples/down-mixes + them itself, so wav and mp3 uploads are forwarded untouched (no decode, no + PCM payload inflation). Other containers (m4a/ogg/webm/flac) are decoded to + a mono WAV. Blocking; call via a thread from async paths. + """ + if b64.startswith("data:"): + b64 = b64.split(",", 1)[1] if "," in b64 else "" + raw = base64.b64decode(b64) + passthrough = _sniff_audio_container(raw) + if passthrough is not None: + return b64, passthrough + + arr, sr = _decode_audio_mono(raw) + arr, sr = _fit_transcoded_audio_to_wav_cap(arr, sr) + return base64.b64encode(_mono_f32_to_wav_bytes(arr, sr)).decode("ascii"), "wav" + + +def _inject_audio_part(messages: list[dict], audio_b64: str, audio_format: str) -> None: + """Append an input_audio part to the last user message, in place. + + Audio rides in the message list like image_url parts do, so it flows through + both the plain and tool-calling generation paths. + """ + part = { + "type": "input_audio", + "input_audio": {"data": audio_b64, "format": audio_format}, + } + for msg in reversed(messages): + if msg.get("role") == "user": + content = msg.get("content") + if isinstance(content, list): + content.append(part) + else: + msg["content"] = [{"type": "text", "text": content or ""}, part] + return + + +def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[str]"]: """ Parse OpenAI-format messages into components the inference backend expects. @@ -1642,24 +2469,22 @@ def _extract_content_parts( (``[{type: "text", ...}, {type: "image_url", ...}]``). Returns: - system_prompt: The system message text (empty string if none provided). + system_prompt: System message text (empty string if none). chat_messages: Non-system messages with content flattened to strings. - image_base64: Base64 data of the *first* image found, or ``None``. + image_base64: Base64 of the *first* image found, or ``None``. """ - system_prompt = "" + system_parts: list[str] = [] chat_messages: list[dict] = [] first_image_b64: Optional[str] = None for msg in messages: - # ── System messages → extract as system_prompt ──────── - if msg.role == "system": + # ── System / developer messages → extract as system_prompt ──────── + if msg.role in ("system", "developer"): if isinstance(msg.content, str): - system_prompt = msg.content + system_parts.append(msg.content) elif isinstance(msg.content, list): # Unlikely but handle: join text parts - system_prompt = "\n".join( - p.text for p in msg.content if p.type == "text" - ) + system_parts.append("\n".join(p.text for p in msg.content if p.type == "text")) continue # ── User / assistant messages ───────────────────────── @@ -1675,29 +2500,27 @@ def _extract_content_parts( elif part.type == "image_url" and first_image_b64 is None: url = part.image_url.url if url.startswith("data:"): - # data:image/png;base64, → extract + # data:image/png;base64, -> extract first_image_b64 = url.split(",", 1)[1] if "," in url else None else: - logger.warning( - f"Remote image URLs not yet supported: {url[:80]}..." - ) + logger.warning(f"Remote image URLs not yet supported: {url[:80]}...") combined_text = "\n".join(text_parts) if text_parts else "" chat_messages.append({"role": msg.role, "content": combined_text}) - return system_prompt, chat_messages, first_image_b64 + return "\n\n".join(p for p in system_parts if p), chat_messages, first_image_b64 # ── External provider proxy ────────────────────────────────────── -# Providers whose stream helper translates `input_document` parts into -# a native attachment block on the wire. For Anthropic the mapping is -# `_stream_anthropic` -> {type:"document", source:...}; for OpenAI it -# is `_stream_openai_responses` -> {type:"input_file", file_data|file_url}. -# Every other provider (gemini / mistral / kimi / openrouter / deepseek / -# custom OpenAI-compat) goes through the generic /chat/completions -# passthrough that forwards messages verbatim, so handing them an -# `input_document` part would 400 with an unknown content_part type. +# Providers whose stream helper translates `input_document` parts into a +# native attachment block on the wire. Anthropic: `_stream_anthropic` -> +# {type:"document", source:...}; OpenAI: `_stream_openai_responses` -> +# {type:"input_file", file_data|file_url}. Every other provider (gemini / +# mistral / kimi / openrouter / deepseek / custom OpenAI-compat) goes through +# the generic /chat/completions passthrough that forwards messages verbatim, +# so handing them an `input_document` part would 400 with an unknown +# content_part type. _INPUT_DOCUMENT_PROVIDERS = frozenset({"anthropic", "openai"}) @@ -1705,6 +2528,7 @@ def _build_external_messages( messages: list, supports_vision: bool, provider_type: Optional[str] = None, + base_url: Optional[str] = None, ) -> list[dict]: """ Convert ChatMessage list to OpenAI-compatible dicts for external providers. @@ -1712,34 +2536,177 @@ def _build_external_messages( Behaviour per content-part type: - `text`: always preserved. - `image_url`: preserved on vision providers; stripped on non-vision. - - `input_document`: preserved ONLY when the provider's stream helper - has explicit translation logic for it (Anthropic + OpenAI today, - see ``_INPUT_DOCUMENT_PROVIDERS``). For every other provider the - part is stripped so the unknown content type doesn't reach generic - /chat/completions passthrough and 400 the request. - - `reasoning`: OpenAI-only Responses reasoning item paired with a - prior tool output. Forwarded ONLY when provider_type=="openai" - so follow-up image edits can replay the required reasoning item. - - `image_generation_call`: OpenAI-only Responses image reference. - Forwarded ONLY when provider_type=="openai" so follow-up image - edits can reference prior generated images. + - `input_document`: preserved ONLY when the provider's stream helper has + explicit translation logic (Anthropic + OpenAI today, see + ``_INPUT_DOCUMENT_PROVIDERS``). Stripped for every other provider so the + unknown type doesn't reach generic /chat/completions and 400. + - `reasoning`: OpenAI-only Responses reasoning item paired with a prior + tool output. Forwarded ONLY when provider_type=="openai" so follow-up + image edits can replay the required reasoning item. + - `image_generation_call`: OpenAI-only Responses image reference. Forwarded + ONLY when provider_type=="openai" so follow-up image edits can reference + prior generated images. - `compaction`: Anthropic-only synthetic part (round-trips server-side compaction state). Forwarded ONLY when provider_type=="anthropic"; - stripped for every other provider so the unknown part doesn't - reach generic /chat/completions passthrough where it would 400 - (e.g. DeepSeek, Mistral, Gemini, Kimi, OpenRouter, etc.). + stripped elsewhere so the unknown part doesn't reach generic + /chat/completions and 400 (DeepSeek, Mistral, Gemini, Kimi, OpenRouter). """ document_provider = provider_type in _INPUT_DOCUMENT_PROVIDERS anthropic = provider_type == "anthropic" openai = provider_type == "openai" + # `extra_content` carries the assistant's text-part `thoughtSignature` + # round-trip on Gemini's native streamGenerateContent endpoint. Custom + # Gemini OpenAI-compat gateways (LiteLLM etc.) route through + # /chat/completions where the field is unknown and can be rejected -- gate + # strictly on the Google-hosted Gemini base. + _native_gemini = False + if provider_type == "gemini" and base_url: + try: + from urllib.parse import urlparse as _urlparse + _host = (_urlparse(base_url).hostname or "").lower() + _native_gemini = _host == "generativelanguage.googleapis.com" + except Exception: + _native_gemini = False + emit_extra_content = _native_gemini + + _SERVER_BUILTIN_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "code_execution", "image_generation"} + ) + + def _is_marked_server_builtin_tool_call(tc: Any) -> bool: + """Return True iff `tc` is a synthetic provider-side tool card with a + canonical builtin name and either: + - the `args._server_tool` marker stamped by the backend, or + - a Gemini `args.google.native_part` payload (durable replay signal + for code_execution / image_generation that predates the marker). + Such cards must not be forwarded to non-native providers: they aren't + real user functions, so the receiving API rejects the orphan tool + history. Real user functions with these names normally have neither + signal. + """ + if not isinstance(tc, dict): + return False + fn = tc.get("function") + if not isinstance(fn, dict): + return False + name = (fn.get("name") or "").lower() + if name not in _SERVER_BUILTIN_TOOL_NAMES: + return False + raw_args = fn.get("arguments") or "" + try: + args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args + except Exception: + return False + if not isinstance(args, dict): + return False + if args.get("_server_tool") is True: + return True + google = args.get("google") + return isinstance(google, dict) and isinstance(google.get("native_part"), dict) + + # When we drop a server-side builtin tool_call, the matching `role="tool"` + # follow-up must also be dropped -- else the provider gets an orphan + # tool_call_id with no matching assistant call, which OpenAI Responses and + # Anthropic both reject. + dropped_server_builtin_tool_call_ids: set[str] = set() + + def _filter_tool_calls(tool_calls: Any) -> Optional[list]: + """Sanitize assistant `tool_calls` for non-native-Gemini providers. + + Two concerns: + 1. `tool_calls[i].extra_content` carries Gemini-only thoughtSignature + metadata; strip it for providers that can't parse the unknown key. + 2. Marked server-side builtin cards (`_server_tool: true` on a + canonical builtin name, or a Gemini `native_part` payload) are + Studio-internal tool cards from a prior native Gemini turn; + forwarding them to OpenAI / Anthropic / custom OAI-compat gateways + sends an orphan `tool_calls` entry (no matching tool declaration, + often no matching `role="tool"` reply) that can be rejected. We + record the dropped call_ids so the matching role=tool message is + skipped below. + Native Gemini keeps both untouched so the translator can replay them + via `native_part`. + """ + if not tool_calls: + return None + if not isinstance(tool_calls, list): + return tool_calls + if emit_extra_content: + return tool_calls + cleaned: list = [] + for _tc in tool_calls: + if _is_marked_server_builtin_tool_call(_tc): + _tc_id = _tc.get("id") if isinstance(_tc, dict) else None + if isinstance(_tc_id, str) and _tc_id: + dropped_server_builtin_tool_call_ids.add(_tc_id) + continue + if not isinstance(_tc, dict): + cleaned.append(_tc) + continue + if "extra_content" not in _tc: + cleaned.append(_tc) + continue + _stripped = {k: v for k, v in _tc.items() if k != "extra_content"} + cleaned.append(_stripped) + return cleaned + result = [] for msg in messages: + # Drop role=tool messages whose matching server-builtin tool_call was + # filtered above. An orphan tool_result with no matching tool_call is + # rejected by OpenAI Responses and Anthropic. + if ( + msg.role == "tool" + and isinstance(msg.tool_call_id, str) + and msg.tool_call_id in dropped_server_builtin_tool_call_ids + ): + continue if isinstance(msg.content, str): - # Skip assistant messages with empty content (some providers reject them) - if msg.role == "assistant" and not msg.content.strip(): + # Drop bare assistant messages with no content AND no tool_calls + # (some providers reject empty assistant turns). Preserve assistant + # turns whose only payload is tool_calls so multi-turn + # function-call loops round-trip. + if msg.role == "assistant" and not msg.content.strip() and not msg.tool_calls: continue - result.append({"role": msg.role, "content": msg.content}) - elif isinstance(msg.content, list): + out: dict[str, Any] = {"role": msg.role, "content": msg.content} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + out["tool_calls"] = _tcs + elif not msg.content.strip(): + # Every tool_call was a dropped synthetic provider card; + # the turn would be an empty + # `{"role":"assistant","content":""}` that some providers + # reject. Skip it entirely. + continue + if msg.role == "tool": + if msg.tool_call_id: + out["tool_call_id"] = msg.tool_call_id + if msg.name: + out["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + out["extra_content"] = msg.extra_content + result.append(out) + continue + # Assistant messages with content=None but populated tool_calls are + # valid (post-tool-call turn). Forward them so the provider helper can + # rebuild the functionCall part. + if msg.content is None and msg.role == "assistant" and msg.tool_calls: + _filtered_tcs = _filter_tool_calls(msg.tool_calls) + if not _filtered_tcs: + # Every tool_call was provider-side synthetic and dropped; + # skip the whole message to avoid an empty assistant turn. + continue + _assistant_only: dict[str, Any] = { + "role": "assistant", + "content": "", + "tool_calls": _filtered_tcs, + } + if emit_extra_content and msg.extra_content: + _assistant_only["extra_content"] = msg.extra_content + result.append(_assistant_only) + continue + if isinstance(msg.content, list): if supports_vision: parts = [] for part in msg.content: @@ -1752,9 +2719,7 @@ def _build_external_messages( "image_url": {"url": part.image_url.url}, } ) - elif ( - part.type == "reasoning" and openai and msg.role == "assistant" - ): + elif part.type == "reasoning" and openai and msg.role == "assistant": reasoning: dict[str, Any] = { "type": "reasoning", "id": part.id, @@ -1764,23 +2729,20 @@ def _build_external_messages( reasoning["status"] = part.status parts.append(reasoning) elif ( - part.type == "image_generation_call" - and openai - and msg.role == "assistant" + part.type == "image_generation_call" and openai and msg.role == "assistant" ): # ExternalProviderClient maps this onto a top-level # Responses input item after the current user prompt, # or onto `previous_response_id` when response_id is - # available from the prior Responses turn. + # available from the prior turn. image_ref = {"type": "image_generation_call", "id": part.id} if getattr(part, "response_id", None): image_ref["response_id"] = part.response_id parts.append(image_ref) elif part.type == "input_document" and document_provider: - # ExternalProviderClient maps this onto - # Anthropic's `document` or OpenAI Responses' - # `input_file` block per provider; every other - # provider would 400 on the unknown part type. + # ExternalProviderClient maps this onto Anthropic's + # `document` or OpenAI Responses' `input_file` block; + # every other provider would 400 on the unknown part. doc: dict[str, Any] = {"type": "input_document"} if part.file_data: doc["file_data"] = part.file_data @@ -1792,20 +2754,35 @@ def _build_external_messages( doc["media_type"] = part.media_type parts.append(doc) elif part.type == "compaction" and anthropic: - # Anthropic stream helper forwards this as a - # native `compaction` block; every other - # provider would 400 on the unknown part, so - # gate by provider_type. + # Anthropic stream helper forwards this as a native + # `compaction` block; every other provider would 400 on + # the unknown part, so gate by provider_type. parts.append({"type": "compaction", "content": part.content}) - if msg.role == "assistant" and not parts: + entry: dict[str, Any] = {"role": msg.role, "content": parts} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + entry["tool_calls"] = _tcs + elif not parts: + # All tool_calls were synthetic and dropped, and no + # content parts survived. Skip rather than forward an + # empty assistant turn that downstream providers reject. + continue + elif msg.role == "assistant" and not parts: continue - result.append({"role": msg.role, "content": parts}) + if msg.role == "tool": + if msg.tool_call_id: + entry["tool_call_id"] = msg.tool_call_id + if msg.name: + entry["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + entry["extra_content"] = msg.extra_content + result.append(entry) else: - # Non-vision provider: strip images / documents, keep - # text, optionally keep compaction (Anthropic only -- + # Non-vision provider: strip images / documents, keep text, + # optionally keep compaction (Anthropic only -- # compaction-capable Anthropic models all report - # supports_vision=True today, but the gate is here for - # safety). + # supports_vision=True today, but gate here for safety). preserved = [] for p in msg.content: if p.type == "text": @@ -1819,11 +2796,7 @@ def _build_external_messages( if p.status: reasoning["status"] = p.status preserved.append(reasoning) - elif ( - p.type == "image_generation_call" - and openai - and msg.role == "assistant" - ): + elif p.type == "image_generation_call" and openai and msg.role == "assistant": image_ref = {"type": "image_generation_call", "id": p.id} if getattr(p, "response_id", None): image_ref["response_id"] = p.response_id @@ -1833,23 +2806,43 @@ def _build_external_messages( if msg.role == "assistant" and not preserved: continue if len(preserved) == 1 and preserved[0]["type"] == "text": - # Single text part collapses back to a string for - # providers that don't accept content arrays. - result.append({"role": msg.role, "content": preserved[0]["text"]}) + # Single text part collapses to a string for providers that + # don't accept content arrays. + entry = {"role": msg.role, "content": preserved[0]["text"]} else: - result.append({"role": msg.role, "content": preserved}) + entry = {"role": msg.role, "content": preserved} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + entry["tool_calls"] = _tcs + else: + # All tool_calls were synthetic and dropped; skip if no + # content survived either. + _entry_content = entry.get("content") + _has_text = ( + isinstance(_entry_content, str) and _entry_content.strip() + ) or (isinstance(_entry_content, list) and len(_entry_content) > 0) + if not _has_text: + continue + if msg.role == "tool": + if msg.tool_call_id: + entry["tool_call_id"] = msg.tool_call_id + if msg.name: + entry["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + entry["extra_content"] = msg.extra_content + result.append(entry) return result async def _proxy_to_external_provider( - payload: ChatCompletionRequest, - request: Request, + payload: ChatCompletionRequest, request: Request ) -> StreamingResponse: """ Proxy a chat completion request to an external LLM provider. - Resolves provider config (from DB or registry), decrypts the API key, - and streams the response back in OpenAI SSE format. + Resolves provider config (DB or registry), decrypts the API key, and + streams the response back in OpenAI SSE format. """ # Resolve provider type and base URL provider_type = payload.provider_type @@ -1903,7 +2896,7 @@ async def _proxy_to_external_provider( detail = "external_model is required when using an external provider.", ) - # Build messages preserving multimodal content for vision-capable providers + # Build messages, preserving multimodal content for vision providers from core.inference.providers import get_provider_info as _get_provider_info _pinfo = _get_provider_info(provider_type) or {} @@ -1912,6 +2905,7 @@ async def _proxy_to_external_provider( payload.messages, _supports_vision, provider_type = provider_type, + base_url = base_url, ) client = ExternalProviderClient( @@ -1920,15 +2914,25 @@ async def _proxy_to_external_provider( api_key = api_key, ) + # `top_k` defaults to 20 in ChatCompletionRequest because the local path + # expects an int, but the external-provider path treats "field omitted from + # JSON" as "use provider default" so callers sending only model/messages + # don't silently get different sampling than before this PR. Pydantic's + # `model_fields_set` tracks explicit-vs-default per request. + _top_k_explicit = payload.top_k if "top_k" in payload.model_fields_set else None + async def _stream(): gen = client.stream_chat_completion( messages = chat_messages, model = model, temperature = payload.temperature, top_p = payload.top_p, - max_tokens = payload.max_tokens, + # Honor max_completion_tokens when max_tokens is absent, so a + # provider-routed request capped only by the newer field still gets + # a limit instead of falling back to the provider default. + max_tokens = _effective_max_tokens(payload), presence_penalty = payload.presence_penalty, - top_k = payload.top_k, + top_k = _top_k_explicit, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, enabled_tools = payload.enabled_tools, @@ -1937,6 +2941,8 @@ async def _proxy_to_external_provider( anthropic_code_exec_container_id = payload.anthropic_code_exec_container_id, prompt_cache_ttl = payload.prompt_cache_ttl, compaction_threshold = payload.compaction_threshold, + tools = payload.tools, + tool_choice = payload.tool_choice, fast_mode = payload.fast_mode, stream = payload.stream, ) @@ -1970,15 +2976,13 @@ async def _proxy_to_external_provider( # ── OpenAI shell-tool container management ─────────────────────── -def _resolve_openai_cloud_client( - body: OpenAIContainerRequest, -) -> ExternalProviderClient: +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. + 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: @@ -2015,9 +3019,7 @@ def _summarize_container(raw: dict) -> OpenAIContainerSummary: 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, + 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, @@ -2031,8 +3033,7 @@ def _summarize_container(raw: dict) -> OpenAIContainerSummary: response_model = ListOpenAIContainersResponse, ) async def list_openai_containers( - body: OpenAIContainerRequest, - current_subject: str = Depends(get_current_subject), + body: OpenAIContainerRequest, current_subject: str = Depends(get_current_subject) ) -> ListOpenAIContainersResponse: """List the user's OpenAI shell-tool containers.""" client = _resolve_openai_cloud_client(body) @@ -2046,13 +3047,16 @@ async def list_openai_containers( detail = f"OpenAI rejected /containers list: {detail}", ) except httpx.HTTPError as exc: - raise HTTPException( - status_code = 502, - detail = f"Failed to reach OpenAI: {exc}", + raise log_and_http_error( + exc, + 502, + "Could not reach OpenAI.", + event = "openai_container_list.transport_error", + log = logger, ) - # 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. + # OpenAI keeps expired containers in /v1/containers indefinitely with + # status="expired" -- dead but still listed. Hide them so the picker + # only shows usable containers. return ListOpenAIContainersResponse( containers = [ _summarize_container(c) @@ -2069,8 +3073,7 @@ async def list_openai_containers( response_model = OpenAIContainerSummary, ) async def create_openai_container( - body: CreateOpenAIContainerBody, - current_subject: str = Depends(get_current_subject), + 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) @@ -2087,9 +3090,12 @@ async def create_openai_container( detail = f"OpenAI rejected /containers create: {detail}", ) except httpx.HTTPError as exc: - raise HTTPException( - status_code = 502, - detail = f"Failed to reach OpenAI: {exc}", + raise log_and_http_error( + exc, + 502, + "Could not reach OpenAI.", + event = "openai_container_create.transport_error", + log = logger, ) if not isinstance(raw, dict): raise HTTPException( @@ -2103,8 +3109,7 @@ async def create_openai_container( @router.post("/external/openai/containers/delete", status_code = 204) async def delete_openai_container( - body: DeleteOpenAIContainerBody, - current_subject: str = Depends(get_current_subject), + body: DeleteOpenAIContainerBody, current_subject: str = Depends(get_current_subject) ) -> None: """Delete a named container by id.""" logger.info( @@ -2134,14 +3139,12 @@ async def delete_openai_container( 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, + raise log_and_http_error( exc, - ) - raise HTTPException( - status_code = 502, - detail = f"Failed to reach OpenAI: {exc}", + 502, + "Could not reach OpenAI.", + event = "openai_container_delete.transport_error", + log = logger, ) finally: await client.close() @@ -2156,35 +3159,80 @@ async def openai_chat_completions( """ OpenAI-compatible chat completions endpoint. - Supports multimodal messages: ``content`` may be a plain string or a - list of content parts (``text`` / ``image_url``). + Supports multimodal messages: ``content`` may be a plain string or a list + of content parts (``text`` / ``image_url``). Non-streaming (default): returns a single ChatCompletion JSON object. Streaming: returns SSE chunks matching OpenAI's format. - ``stream`` defaults to ``false`` to match OpenAI's spec; clients opt - into SSE by sending ``stream: true``. + ``stream`` defaults to ``false`` per OpenAI's spec; clients opt into SSE by + sending ``stream: true``. - Automatically routes to the correct backend: + Routes to the correct backend automatically: - GGUF models → llama-server via LlamaCppBackend - Other models → Unsloth/transformers via InferenceBackend """ + # OpenAI's newer "developer" role is equivalent to "system". Normalize it + # before provider routing so external providers (which may not accept the + # "developer" role) get "system" too, matching the local path. + for _m in payload.messages: + if _m.role == "developer": + _m.role = "system" + + if payload.logprobs: + _raise_unsupported_openai_parameter( + "logprobs", "logprobs is not supported for chat completions." + ) + if payload.top_logprobs is not None: + _raise_unsupported_openai_parameter( + "top_logprobs", "top_logprobs is not supported for chat completions." + ) + # ── External provider routing ──────────────────────────────── - # encrypted_api_key is optional — local providers (llama.cpp / vLLM / Ollama) may run without auth. + # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + if _wants_multiple_choices(payload): + _raise_unsupported_n("external provider chat completions") return await _proxy_to_external_provider(payload, request) + # Reject a malformed function tool here: it would otherwise reach + # llama-server and surface as an opaque 500 "Failed to parse tools". + if payload.tools: + for _tool in payload.tools: + if not isinstance(_tool, dict): + continue + # llama-server 500s ("Failed to parse tools: Missing tool type") when + # a function tool omits "type". Default it to "function" so a + # well-formed tool isn't rejected over a missing discriminator (and a + # malformed one still surfaces as a clean 400 below, not a 500). + if _tool.get("type") is None and isinstance(_tool.get("function"), dict): + _tool["type"] = "function" + if _tool.get("type") != "function": + continue + _fn = _tool.get("function") + _name = _fn.get("name") if isinstance(_fn, dict) else None + if not isinstance(_name, str) or not _name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tools': each tool must have a 'function' with a 'name'.", + status = 400, + code = "invalid_value", + param = "tools", + ), + ) + llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded - # OpenAI-SDK clients send ``chat_template_kwargs`` via ``extra_body``, - # which the SDK spreads into the request body at the top level. Studio's + # OpenAI-SDK clients send ``chat_template_kwargs`` via ``extra_body``, which + # the SDK spreads into the request body at the top level. Studio's # ChatCompletionRequest has ``extra="allow"`` so pydantic stashes them in - # ``model_extra``, but the typed ``payload.enable_thinking`` path is what - # downstream generators actually consume. Lift ``enable_thinking`` from - # the extra-body chat_template_kwargs onto the typed field so clients - # that only know the OpenAI shape (data_designer recipe runs, etc.) - # can still control the reasoning preamble. + # ``model_extra``, but downstream generators consume the typed + # ``payload.enable_thinking``. Lift ``enable_thinking`` from the extra-body + # chat_template_kwargs onto the typed field so clients that only know the + # OpenAI shape (data_designer recipe runs, etc.) can still control the + # reasoning preamble. _extra = getattr(payload, "model_extra", None) if payload.enable_thinking is None and isinstance(_extra, dict): _tpl_kw = _extra.get("chat_template_kwargs") @@ -2192,9 +3240,13 @@ async def openai_chat_completions( payload.enable_thinking = bool(_tpl_kw["enable_thinking"]) # ── Determine which backend is active ───────────────────── + # Single-model server: any model name serves the loaded model (drop-in + # OpenAI compat), so payload.model is only a fallback label here. if using_gguf: model_name = llama_backend.model_identifier or payload.model if getattr(llama_backend, "_is_audio", False): + if _wants_multiple_choices(payload): + _raise_unsupported_n("GGUF audio chat completions") return await generate_audio(payload, request) else: backend = get_inference_backend() @@ -2204,9 +3256,11 @@ async def openai_chat_completions( detail = "No model loaded. Call POST /inference/load first.", ) model_name = backend.active_model_name or payload.model + if _wants_multiple_choices(payload): + _raise_unsupported_n("non-GGUF chat completions") # ── Audio TTS path: auto-route to audio generation ──── - # (Whisper is ASR not TTS — handled below in audio input path) + # (Whisper is ASR not TTS -- handled below in audio input path) model_info = backend.models.get(backend.active_model_name, {}) if model_info.get("is_audio") and model_info.get("audio_type") != "whisper": return await generate_audio(payload, request) @@ -2240,7 +3294,7 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, + max_new_tokens = _effective_max_tokens(payload) or 2048, repetition_penalty = payload.repetition_penalty, cancel_event = cancel_event, ) @@ -2294,9 +3348,7 @@ async def openai_chat_completions( id = completion_id, created = created, model = model_name, - choices = [ - ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop") - ], + choices = [ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" @@ -2304,9 +3356,7 @@ async def openai_chat_completions( cancel_event.set() raise except Exception as e: - logger.error( - f"Error during audio input streaming: {e}", exc_info = True - ) + logger.error(f"Error during audio input streaming: {e}", exc_info = True) yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" finally: _tracker.__exit__(None, None, None) @@ -2344,15 +3394,18 @@ async def openai_chat_completions( # unaware of `role="tool"` messages and assistant messages that only # carry `tool_calls` (content=None) — both of which are valid in # multi-turn client-side tool loops. + effective_max_tokens = _effective_max_tokens(payload) + + normalized_stop = _normalize_stop_sequences(payload.stop) + _has_tool_messages = any(m.role == "tool" or m.tool_calls for m in payload.messages) # Route guided-decoding requests through the verbatim passthrough so - # ``response_format`` (JSON schema) actually reaches llama-server and - # the model's GBNF-constrained output comes back unmodified. The - # non-passthrough GGUF path below calls ``generate_chat_completion`` - # which has no response_format kwarg, so the schema gets silently - # dropped and data_designer falls back to free-form sampling. Guided - # decoding does not require ``supports_tools`` - the grammar machinery - # is independent of tool-call parsing. + # ``response_format`` (JSON schema) reaches llama-server and the model's + # GBNF-constrained output comes back unmodified. The non-passthrough GGUF + # path below calls ``generate_chat_completion`` which has no response_format + # kwarg, so the schema gets silently dropped and data_designer falls back to + # free-form sampling. Guided decoding does not require ``supports_tools`` -- + # the grammar machinery is independent of tool-call parsing. _has_response_format = _extract_response_format(payload) is not None _tools_passthrough = llama_backend.supports_tools and ( (payload.tools and len(payload.tools) > 0) or _has_tool_messages @@ -2362,16 +3415,21 @@ async def openai_chat_completions( and not _effective_enable_tools(payload) and (_tools_passthrough or _has_response_format) ): + if _wants_multiple_choices(payload): + _raise_unsupported_n("GGUF tool or response_format passthrough") if payload.audio_base64: + # This path forwards the request verbatim, so the transcoded audio + # never gets injected. (The agentic tool loop below does support + # audio.) raise HTTPException( status_code = 400, - detail = "Audio input is not supported for GGUF chat models yet.", + detail = "Audio input is not supported together with guided decoding or client-supplied tools yet.", ) - # Preserve the vision guard that would otherwise run in the - # non-passthrough path below: text-only tool-capable GGUFs - # should return a clear 400 here rather than forwarding the - # image to llama-server and surfacing an opaque upstream error. + # Preserve the vision guard from the non-passthrough path below: + # text-only tool-capable GGUFs should return a clear 400 here rather + # than forwarding the image to llama-server and surfacing an opaque + # upstream error. if not llama_backend.is_vision and ( payload.image_base64 or any( @@ -2388,8 +3446,8 @@ async def openai_chat_completions( cancel_event = threading.Event() completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" # `stream` defaults to False on ChatCompletionRequest (OpenAI spec - # parity). Naive curl / .NET / System.Text.Json clients omitting - # the field used to get SSE here and choke on deserialization (#5047). + # parity). Naive curl / .NET / System.Text.Json clients omitting the + # field used to get SSE here and choke on deserialization (#5047). if payload.stream: return await _openai_passthrough_stream( request, @@ -2406,9 +3464,7 @@ async def openai_chat_completions( ) # ── Parse messages (handles multimodal content parts) ───── - system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts( - payload.messages - ) + system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) if not chat_messages: raise HTTPException( @@ -2418,17 +3474,43 @@ async def openai_chat_completions( # ── GGUF path: proxy to llama-server /v1/chat/completions ── if using_gguf: + # Forward uploaded audio as an input_audio part. wav/mp3 pass through + # untouched (llama-server decodes and resamples them via the mmproj + # audio encoder); other containers are transcoded to WAV here. The part + # is injected into the message list below so it rides through both the + # plain and tool-calling paths, exactly like image_url parts. + audio_b64 = None + audio_format = "wav" if payload.audio_base64: - raise HTTPException( - status_code = 400, - detail = "Audio input is not supported for GGUF chat models yet.", - ) + if not getattr(llama_backend, "_has_audio_input", False): + raise HTTPException( + status_code = 400, + detail = "Audio provided but current GGUF model does not support audio input.", + ) + if len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: + raise HTTPException( + status_code = 413, + detail = "Audio file is too large (max ~25 MB).", + ) + try: + audio_b64, audio_format = await asyncio.to_thread( + _prepare_audio_for_llama, payload.audio_base64 + ) + except Exception as e: + logger.warning("Audio decode failed: %s", e, exc_info = True) + raise HTTPException( + status_code = 400, + detail = "Could not decode the provided audio file.", + ) - gguf_messages, has_gguf_image = _openai_messages_for_gguf_chat( + gguf_messages, _ = _openai_messages_for_gguf_chat( payload, llama_backend.is_vision, ) + gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt) image_b64 = None + if audio_b64: + _inject_audio_part(gguf_messages, audio_b64, audio_format) cancel_event = threading.Event() @@ -2437,96 +3519,91 @@ async def openai_chat_completions( # ── Tool-calling path (agentic loop) ────────────────── # `_effective_enable_tools` lets `unsloth run --enable-tools/--disable-tools` - # hard-override the per-request value. Without a CLI override, falls - # back to `payload.enable_tools` (existing behavior). - use_tools = ( - _effective_enable_tools(payload) - and llama_backend.supports_tools - and not has_gguf_image - ) + # hard-override the per-request value, else falls back to + # `payload.enable_tools`. `mcp_enabled=true` also opens the tool loop so + # MCP-only callers needn't flip a second flag, BUT must still honor a + # CLI `--disable-tools` policy -- checking the raw policy here keeps + # `mcp_enabled` from re-enabling tools the operator explicitly forbade. + from state.tool_policy import get_tool_policy as _get_tool_policy_g + + _cli_policy = _get_tool_policy_g() + _tools_on = _effective_enable_tools(payload) + _mcp_allowed = bool(payload.mcp_enabled) and _cli_policy is not False + use_tools = (_tools_on or _mcp_allowed) and llama_backend.supports_tools if use_tools: - from core.inference.tools import ALL_TOOLS + from core.inference.tools import ALL_TOOLS, get_enabled_mcp_tools - if payload.enabled_tools is not None: + if not _tools_on: + # MCP-only request: skip built-ins, leave room for MCP tools. + tools_to_use = [] + elif payload.enabled_tools is not None: tools_to_use = [ - t - for t in ALL_TOOLS - if t["function"]["name"] in payload.enabled_tools + t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools ] else: tools_to_use = ALL_TOOLS + # Drop the RAG tool without a scope: nothing to search over. + if not payload.rag_scope: + tools_to_use = [ + t for t in tools_to_use if t["function"]["name"] != "search_knowledge_base" + ] + + if _mcp_allowed: + tools_to_use = tools_to_use + await get_enabled_mcp_tools() + + # Skip the tool loop when no tool survived, so the safetensors + # loop's "empty = allow all" semantic can't reach built-in tools + # the caller didn't opt into. Callers who omit enabled_tools still + # get ALL_TOOLS here, so this only suppresses the loop when + # discovery + opt-in left it genuinely empty. + if not tools_to_use: + use_tools = False + + if use_tools: + if _wants_multiple_choices(payload): + _raise_unsupported_n("GGUF tool chat completions") # ── Tool-use system prompt nudge ────────────────────── - _tool_names = {t["function"]["name"] for t in tools_to_use} - _has_web = "web_search" in _tool_names - _has_code = "python" in _tool_names or "terminal" in _tool_names - - _date_line = f"The current date is {_date.today().isoformat()}." - - # Small models (<9B) struggle with multi-step search plans, - # so simplify the web tips to avoid plan-then-stall behavior. - _model_size_b = _extract_model_size_b(model_name) - _is_small_model = _model_size_b is not None and _model_size_b < 9 - - if _is_small_model: - _web_tips = "Do not repeat the same search query." - else: - _web_tips = ( - "When you search and find a relevant URL in the results, " - "fetch its full content by calling web_search with the url parameter. " - "Do not repeat the same search query. If a search returns " - "no useful results, try rephrasing or fetching a result URL directly." - ) - _code_tips = ( - "Use code execution for math, calculations, data processing, " - "or to parse and analyze information from tool results." + _nudge = _build_tool_action_nudge( + tools = tools_to_use, + model_name = model_name, ) - if _has_web and _has_code: - _nudge = ( - _date_line + " " - "You have access to tools. When appropriate, prefer using " - "tools rather than answering from memory. " - + _web_tips - + " " - + _code_tips + # Nudge the model to ground in attached documents instead of memory. + _tool_names = {(t.get("function") or {}).get("name") for t in (tools_to_use or [])} + _rag_active = "search_knowledge_base" in _tool_names and payload.rag_scope + if _rag_active: + _rag_nudge = ( + "The user has attached documents to this conversation. Relevant " + "passages are retrieved and provided to you automatically; base " + "your answer on them and cite them. You can also call " + "search_knowledge_base to look for more. Do not answer from " + "memory when the attached documents are relevant." ) - elif _has_code: - _nudge = ( - _date_line + " " - "You have access to tools. When appropriate, prefer using " - "code execution rather than answering from memory. " + _code_tips - ) - elif _has_web: - _nudge = ( - _date_line + " " - "You have access to tools. When appropriate, prefer using " - "web search for up-to-date or uncertain factual " - "information rather than answering from memory. " + _web_tips - ) - else: - _nudge = "" + # Prefix the date when the tool nudge is empty (RAG-only tool set). + _date_line = f"The current date is {_date.today().isoformat()}." + _nudge = _date_line + " " + _rag_nudge if not _nudge else _nudge + " " + _rag_nudge if _nudge: - _nudge += _TOOL_ACTION_NUDGE # Append nudge to system prompt (preserve user's prompt) if system_prompt: system_prompt = system_prompt.rstrip() + "\n\n" + _nudge else: system_prompt = _nudge - # Rebuild gguf_messages with updated system prompt - gguf_messages = [] - if system_prompt: - gguf_messages.append({"role": "system", "content": system_prompt}) - gguf_messages.extend(chat_messages) + gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt) + + _gguf_auto_heal_tool_calls = ( + payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True + ) # ── Strip stale tool-call XML from conversation history ─ for _msg in gguf_messages: - if _msg.get("role") == "assistant" and isinstance( - _msg.get("content"), str - ): - _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + _msg["content"] = _strip_tool_xml_for_display( + _msg["content"], + auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + ).strip() def gguf_generate_with_tools(): return llama_backend.generate_chat_completion_with_tools( @@ -2536,16 +3613,16 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_tokens = payload.max_tokens, + max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, presence_penalty = payload.presence_penalty, + stop = normalized_stop, cancel_event = cancel_event, + seed = payload.seed, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, - auto_heal_tool_calls = payload.auto_heal_tool_calls - if payload.auto_heal_tool_calls is not None - else True, + auto_heal_tool_calls = _gguf_auto_heal_tool_calls, max_tool_iterations = payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25, @@ -2553,6 +3630,8 @@ async def openai_chat_completions( if payload.tool_call_timeout is not None else 300, session_id = payload.session_id, + rag_scope = payload.rag_scope, + disable_parallel_tool_use = payload.parallel_tool_calls is False, ) _tool_sentinel = object() @@ -2576,12 +3655,13 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - # Iterate the synchronous generator in a thread so - # the event loop stays free for disconnect detection. + # Iterate the sync generator in a thread so the event loop + # stays free for disconnect detection. gen = gguf_generate_with_tools() prev_text = "" _stream_usage = None _stream_timings = None + _stream_finish = None while True: if cancel_event.is_set(): break @@ -2594,14 +3674,14 @@ async def openai_chat_completions( break if event["type"] == "status": - # Empty status marks an iteration boundary - # in the GGUF tool loop (e.g. after a - # re-prompt). Reset the cumulative cursor - # so the next assistant turn streams cleanly. + # Empty status marks an iteration boundary in the + # GGUF tool loop (e.g. after a re-prompt). Reset the + # cumulative cursor so the next assistant turn + # streams cleanly. if not event["text"]: prev_text = "" - # Emit tool status as a custom SSE event - # (including empty ones to clear UI badges) + # Emit tool status as a custom SSE event (including + # empty ones to clear UI badges) status_data = json.dumps( { "type": "tool_status", @@ -2620,14 +3700,17 @@ async def openai_chat_completions( if event["type"] == "metadata": _stream_usage = event.get("usage") _stream_timings = event.get("timings") + _stream_finish = event.get("finish_reason") continue - # "content" type -- cumulative text - # Sanitize the full cumulative then diff against - # the last sanitized snapshot so cross-chunk XML - # tags are handled correctly. + # "content" type -- cumulative text. Sanitize the full + # cumulative then diff against the last sanitized + # snapshot so cross-chunk XML tags are handled correctly. raw_cumulative = event.get("text", "") - clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative) + clean_cumulative = _strip_tool_xml_for_display( + raw_cumulative, + auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative if not new_text: @@ -2652,29 +3735,21 @@ async def openai_chat_completions( choices = [ ChunkChoice( delta = ChoiceDelta(), - finish_reason = "stop", + finish_reason = _clamp_finish_reason(_stream_finish), ) ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" - # Usage chunk (OpenAI-standard: choices=[], usage populated) - if _stream_usage or _stream_timings: - usage_obj = CompletionUsage( - prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0), - completion_tokens = (_stream_usage or {}).get( - "completion_tokens", 0 - ), - total_tokens = (_stream_usage or {}).get("total_tokens", 0), - ) - usage_chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [], - usage = usage_obj, - timings = _stream_timings, - ) - yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" + usage_line = _openai_stream_usage_chunk( + payload, + completion_id, + created, + model_name, + _stream_usage, + _stream_timings, + ) + if usage_line is not None: + yield usage_line yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -2685,12 +3760,7 @@ async def openai_chat_completions( tb = traceback.format_exc() logger.error(f"Error during GGUF tool streaming: {e}\n{tb}") - error_chunk = { - "error": { - "message": _friendly_error(e), - "type": "server_error", - }, - } + error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: _tracker.__exit__(None, None, None) @@ -2707,7 +3777,10 @@ async def openai_chat_completions( # ── Standard GGUF path (no tools) ───────────────────── - def gguf_generate(): + def gguf_generate(choice_index: int = 0): + _seed = payload.seed + if _seed is not None and _seed >= 0 and choice_index: + _seed += choice_index return llama_backend.generate_chat_completion( messages = gguf_messages, image_b64 = image_b64, @@ -2715,18 +3788,22 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_tokens = payload.max_tokens, + max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, presence_penalty = payload.presence_penalty, + stop = normalized_stop, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, + seed = _seed, ) _gguf_sentinel = object() if payload.stream: + if _wants_multiple_choices(payload): + _raise_unsupported_n("streaming GGUF chat completions") _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() @@ -2747,12 +3824,13 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - # Iterate the synchronous generator in a thread so - # the event loop stays free for disconnect detection. + # Iterate the sync generator in a thread so the event loop + # stays free for disconnect detection. gen = gguf_generate() prev_text = "" _stream_usage = None _stream_timings = None + _stream_finish = None while True: if cancel_event.is_set(): break @@ -2762,19 +3840,16 @@ async def openai_chat_completions( cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) if cumulative is _gguf_sentinel: break - # Capture server metadata for final usage chunk + # Capture server metadata for the final usage chunk if isinstance(cumulative, dict): if cumulative.get("type") == "metadata": _stream_usage = cumulative.get("usage") _stream_timings = cumulative.get("timings") + _stream_finish = cumulative.get("finish_reason") else: logger.warning( "gguf_stream_chunks: unexpected dict event: %s", - { - k: v - for k, v in cumulative.items() - if k != "timings" - }, + {k: v for k, v in cumulative.items() if k != "timings"}, ) continue new_text = cumulative[len(prev_text) :] @@ -2802,29 +3877,21 @@ async def openai_chat_completions( choices = [ ChunkChoice( delta = ChoiceDelta(), - finish_reason = "stop", + finish_reason = _clamp_finish_reason(_stream_finish), ) ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" - # Usage chunk (OpenAI-standard: choices=[], usage populated) - if _stream_usage or _stream_timings: - usage_obj = CompletionUsage( - prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0), - completion_tokens = (_stream_usage or {}).get( - "completion_tokens", 0 - ), - total_tokens = (_stream_usage or {}).get("total_tokens", 0), - ) - usage_chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [], - usage = usage_obj, - timings = _stream_timings, - ) - yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" + usage_line = _openai_stream_usage_chunk( + payload, + completion_id, + created, + model_name, + _stream_usage, + _stream_timings, + ) + if usage_line is not None: + yield usage_line yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -2832,12 +3899,7 @@ async def openai_chat_completions( raise except Exception as e: logger.error(f"Error during GGUF streaming: {e}", exc_info = True) - error_chunk = { - "error": { - "message": _friendly_error(e), - "type": "server_error", - }, - } + error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: _tracker.__exit__(None, None, None) @@ -2853,28 +3915,75 @@ async def openai_chat_completions( ) else: try: - full_text = "" - for token in gguf_generate(): - if isinstance(token, dict): - continue # skip metadata dict in non-streaming path - full_text = token + # ``n`` requests several independent completions; the single + # decode slot yields one at a time, so loop sequentially. + _n = payload.n or 1 + + _choices = [] + _prompt_tokens = 0 + _sum_completion = 0 + _prompt_details = None + for _idx in range(_n): + # Stop spawning the remaining choices once cancelled. + if cancel_event.is_set(): + break + full_text = "" + completion_usage = None + completion_finish = None + for token in gguf_generate(_idx): + if isinstance(token, dict): + if token.get("type") == "metadata": + completion_usage = token.get("usage") + completion_finish = token.get("finish_reason") + continue + full_text = token + + _choices.append( + CompletionChoice( + index = _idx, + message = CompletionMessage(content = full_text), + finish_reason = _clamp_finish_reason(completion_finish), + ) + ) + if completion_usage: + # The prompt is shared across all n choices, so count its + # tokens ONCE (OpenAI bills only generated tokens for each + # extra choice). Only completion_tokens accumulates. + _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens + _sum_completion += completion_usage.get("completion_tokens") or 0 + if _prompt_details is None: + _prompt_details = completion_usage.get("prompt_tokens_details") response = ChatCompletion( id = completion_id, created = created, model = model_name, - choices = [ - CompletionChoice( - message = CompletionMessage(content = full_text), - finish_reason = "stop", - ) - ], + choices = _choices, + usage = CompletionUsage( + prompt_tokens = _prompt_tokens, + completion_tokens = _sum_completion, + total_tokens = _prompt_tokens + _sum_completion, + prompt_tokens_details = _prompt_tokens_details(_prompt_details), + ), ) return JSONResponse(content = response.model_dump()) except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) - raise HTTPException(status_code = 500, detail = str(e)) + # An over-context prompt makes llama-server return 400; map any + # upstream 4xx to a 400 client error rather than leaking a 500. + _cls = _classify_llama_generation_error(e) + if _cls is not None: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + _friendly_error(e), + status = 400, + code = "context_length_exceeded" if _cls else None, + param = "messages", + ), + ) + raise HTTPException(status_code = 500, detail = safe_error_detail(e)) # ── Standard Unsloth path ───────────────────────────────── @@ -2902,7 +4011,13 @@ async def openai_chat_completions( except HTTPException: raise except Exception as e: - raise HTTPException(status_code = 400, detail = f"Failed to decode image: {e}") + raise log_and_http_error( + e, + 400, + "Failed to decode image", + event = "inference.decode_image_failed", + log = logger, + ) # Classify capability flags from the loaded template. _sf_model_info = backend.models.get(backend.active_model_name, {}) @@ -2914,26 +4029,29 @@ async def openai_chat_completions( created = int(time.time()) # ── Safetensors tool-calling path ───────────────────────── - # Mirrors the GGUF agentic loop's event shape. Disabled for - # vision turns (untested overlap with image render slot) and - # for gpt-oss (Harmony uses dedicated channels, not - # XML -- gpt-oss tools still work via the GGUF path). + # Mirrors the GGUF agentic loop's event shape. Disabled for vision turns + # (untested overlap with image render slot) and for gpt-oss (Harmony uses + # dedicated channels, not XML -- gpt-oss tools still work via + # the GGUF path). _sf_is_gptoss = False try: - _sf_is_gptoss = bool( - hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model() - ) + _sf_is_gptoss = bool(hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model()) except Exception: _sf_is_gptoss = False _sf_tool_budget = ( - payload.max_tool_calls_per_message - if payload.max_tool_calls_per_message is not None - else 25 + payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25 ) + # Match the GGUF path: mcp_enabled also opens the tool loop on its own + # but must still honor a CLI `--disable-tools` policy. + from state.tool_policy import get_tool_policy as _get_tool_policy_sf + + _sf_cli_policy = _get_tool_policy_sf() + _sf_tools_on = _effective_enable_tools(payload) + _sf_mcp_allowed = bool(payload.mcp_enabled) and _sf_cli_policy is not False _sf_use_tools = ( - _effective_enable_tools(payload) + (_sf_tools_on or _sf_mcp_allowed) and _sf_features.get("supports_tools", False) and image is None and not _sf_is_gptoss @@ -2941,70 +4059,68 @@ async def openai_chat_completions( ) if _sf_use_tools: - from core.inference.tools import ALL_TOOLS + from core.inference.tools import ALL_TOOLS, get_enabled_mcp_tools - if payload.enabled_tools is not None: + if not _sf_tools_on: + _sf_tools_to_use = [] + elif payload.enabled_tools is not None: _sf_tools_to_use = [ t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools ] else: _sf_tools_to_use = ALL_TOOLS - _sf_tool_names = {t["function"]["name"] for t in _sf_tools_to_use} - _sf_has_web = "web_search" in _sf_tool_names - _sf_has_code = "python" in _sf_tool_names or "terminal" in _sf_tool_names + # Drop the RAG tool unless the request carries a retrieval scope. + if not payload.rag_scope: + _sf_tools_to_use = [ + t for t in _sf_tools_to_use if t["function"]["name"] != "search_knowledge_base" + ] - _sf_date_line = f"The current date is {_date.today().isoformat()}." - _sf_model_size_b = _extract_model_size_b(model_name) - _sf_is_small_model = _sf_model_size_b is not None and _sf_model_size_b < 9 + if _sf_mcp_allowed: + _sf_tools_to_use = _sf_tools_to_use + await get_enabled_mcp_tools() - if _sf_is_small_model: - _sf_web_tips = "Do not repeat the same search query." - else: - _sf_web_tips = ( - "When you search and find a relevant URL in the results, " - "fetch its full content by calling web_search with the url parameter. " - "Do not repeat the same search query. If a search returns " - "no useful results, try rephrasing or fetching a result URL directly." - ) - _sf_code_tips = ( - "Use code execution for math, calculations, data processing, " - "or to parse and analyze information from tool results." + # Mirror the GGUF path: refuse to enter the tool loop when nothing + # survived, so a model-emitted built-in call can't piggy-back on the + # empty allow-list. + if not _sf_tools_to_use: + _sf_use_tools = False + + if _sf_use_tools: + _sf_nudge = _build_tool_action_nudge( + tools = _sf_tools_to_use, + model_name = model_name, ) - if _sf_has_web and _sf_has_code: - _sf_nudge = ( - _sf_date_line + " " - "You have access to tools. When appropriate, prefer using " - "tools rather than answering from memory. " - + _sf_web_tips - + " " - + _sf_code_tips + # RAG nudge, mirroring the GGUF path. + _sf_tool_names = {(t.get("function") or {}).get("name") for t in (_sf_tools_to_use or [])} + _sf_rag_active = "search_knowledge_base" in _sf_tool_names and payload.rag_scope + if _sf_rag_active: + _sf_rag_nudge = ( + "The user has attached documents to this conversation. Relevant " + "passages are retrieved and provided to you automatically; base " + "your answer on them and cite them. You can also call " + "search_knowledge_base to look for more. Do not answer from " + "memory when the attached documents are relevant." ) - elif _sf_has_code: + # Prefix the date when the tool nudge is empty (RAG-only tool set). + _sf_date_line = f"The current date is {_date.today().isoformat()}." _sf_nudge = ( - _sf_date_line + " " - "You have access to tools. When appropriate, prefer using " - "code execution rather than answering from memory. " + _sf_code_tips + _sf_date_line + " " + _sf_rag_nudge + if not _sf_nudge + else _sf_nudge + " " + _sf_rag_nudge ) - elif _sf_has_web: - _sf_nudge = ( - _sf_date_line + " " - "You have access to tools. When appropriate, prefer using " - "web search for up-to-date or uncertain factual " - "information rather than answering from memory. " + _sf_web_tips - ) - else: - _sf_nudge = "" _sf_system_prompt = system_prompt if _sf_nudge: - _sf_nudge += _TOOL_ACTION_NUDGE if _sf_system_prompt: _sf_system_prompt = _sf_system_prompt.rstrip() + "\n\n" + _sf_nudge else: _sf_system_prompt = _sf_nudge + _sf_auto_heal_tool_calls = ( + payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True + ) + # Strip stale tool-call XML from prior assistant turns. _sf_chat_messages = [] for _msg in chat_messages: @@ -3012,12 +4128,18 @@ async def openai_chat_completions( _sf_chat_messages.append( { **_msg, - "content": _TOOL_XML_RE.sub("", _msg["content"]).strip(), + "content": _strip_tool_xml_for_display( + _msg["content"], + auto_heal_tool_calls = _sf_auto_heal_tool_calls, + ).strip(), } ) else: _sf_chat_messages.append(_msg) + # Request-scoped usage/timings receptacle (filled at gen_done). + _sf_stats_holder: dict = {} + def sf_generate_with_tools(): return backend.generate_chat_completion_with_tools( messages = _sf_chat_messages, @@ -3027,21 +4149,21 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_tokens = payload.max_tokens, + max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, - auto_heal_tool_calls = payload.auto_heal_tool_calls - if payload.auto_heal_tool_calls is not None - else True, + auto_heal_tool_calls = _sf_auto_heal_tool_calls, max_tool_iterations = _sf_tool_budget, tool_call_timeout = payload.tool_call_timeout if payload.tool_call_timeout is not None else 300, session_id = payload.session_id, + rag_scope = payload.rag_scope, use_adapter = payload.use_adapter, + stats_holder = _sf_stats_holder, ) _sf_tool_sentinel = object() @@ -3099,7 +4221,10 @@ async def openai_chat_completions( # Diff cumulative cleaned text against last snapshot. raw_cumulative = event.get("text", "") - clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative) + clean_cumulative = _strip_tool_xml_for_display( + raw_cumulative, + auto_heal_tool_calls = _sf_auto_heal_tool_calls, + ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative if not new_text: @@ -3129,6 +4254,21 @@ async def openai_chat_completions( ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" + # Usage chunk from the last turn, same shape as the + # GGUF tool loop's metadata. Request-scoped holder, so + # concurrent streams cannot read each other's stats. + _stats = _sf_stats_holder.get("stats") + if _stats: + usage_line = _openai_stream_usage_chunk( + payload, + completion_id, + created, + model_name, + _stats.get("usage"), + _stats.get("timings"), + ) + if usage_line is not None: + yield usage_line yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -3137,8 +4277,8 @@ async def openai_chat_completions( raise except Exception: backend.reset_generation_state() - # Generic wire message; full trace stays in the log - # (CWE-209: transformers/torch errors may leak paths). + # Generic wire message; full trace stays in the log (CWE-209: + # transformers/torch errors may leak paths). logger.exception("safetensors tool stream error") error_chunk = { "error": { @@ -3171,7 +4311,10 @@ async def openai_chat_completions( if cancel_event.is_set(): break if event.get("type") == "content": - full_text = _TOOL_XML_RE.sub("", event.get("text", "")) + full_text = _strip_tool_xml_for_display( + event.get("text", ""), + auto_heal_tool_calls = _sf_auto_heal_tool_calls, + ) return full_text content_text = await asyncio.to_thread(_drain_to_text) @@ -3207,11 +4350,11 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, + max_new_tokens = effective_max_tokens or 2048, repetition_penalty = payload.repetition_penalty, ) - # Forward reasoning kwargs; the worker/template wrapper peels off - # any the template doesn't accept. + # Forward reasoning kwargs; the worker/template wrapper peels off any the + # template doesn't accept. if payload.enable_thinking is not None: gen_kwargs["enable_thinking"] = payload.enable_thinking if payload.reasoning_effort is not None: @@ -3219,19 +4362,25 @@ async def openai_chat_completions( if payload.preserve_thinking is not None: gen_kwargs["preserve_thinking"] = payload.preserve_thinking + # Request-scoped usage/timings receptacle (filled at gen_done). + stats_holder: dict = {} + if payload.use_adapter is not None: def generate(): return backend.generate_with_adapter_control( use_adapter = payload.use_adapter, cancel_event = cancel_event, + stats_holder = stats_holder, **gen_kwargs, ) else: def generate(): return backend.generate_chat_response( - cancel_event = cancel_event, **gen_kwargs + cancel_event = cancel_event, + stats_holder = stats_holder, + **gen_kwargs, ) # ── Streaming response ──────────────────────────────────────── @@ -3256,12 +4405,12 @@ async def openai_chat_completions( yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" prev_text = "" - # Run sync generator in thread pool to avoid blocking - # the event loop. Critical for compare mode: two SSE - # requests arrive concurrently but the orchestrator - # serializes them via _gen_lock. Without run_in_executor - # the second request's blocking lock acquisition would - # freeze the entire event loop, stalling both streams. + # Run the sync generator in a thread pool to avoid blocking the + # event loop. Critical for compare mode: two SSE requests arrive + # concurrently but the orchestrator serializes them via + # _gen_lock; without run_in_executor the second request's + # blocking lock acquisition would freeze the entire event loop, + # stalling both streams. _DONE = object() # sentinel for generator exhaustion loop = asyncio.get_event_loop() gen = generate() @@ -3270,8 +4419,8 @@ async def openai_chat_completions( backend.reset_generation_state() break # next(gen, _DONE) returns _DONE instead of raising - # StopIteration — StopIteration cannot propagate - # through asyncio futures (Python limitation). + # StopIteration -- StopIteration can't propagate through + # asyncio futures (Python limitation). cumulative = await loop.run_in_executor(None, next, gen, _DONE) if cumulative is _DONE: break @@ -3308,6 +4457,22 @@ async def openai_chat_completions( ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" + # Usage chunk (choices=[], usage set), same shape as the + # GGUF path so the speed popover works for MLX too. + # Request-scoped holder, so concurrent streams cannot + # read each other's stats. + _stats = stats_holder.get("stats") + if _stats: + usage_line = _openai_stream_usage_chunk( + payload, + completion_id, + created, + model_name, + _stats.get("usage"), + _stats.get("timings"), + ) + if usage_line is not None: + yield usage_line yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -3360,7 +4525,7 @@ async def openai_chat_completions( except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI completion: {e}", exc_info = True) - raise HTTPException(status_code = 500, detail = str(e)) + raise HTTPException(status_code = 500, detail = safe_error_detail(e)) # ===================================================================== @@ -3387,8 +4552,8 @@ async def serve_sandbox_file( """ Serve image files created by Python tool execution. - Accepts auth via Authorization header OR ?token= query param - (needed because cannot send custom headers). + Accepts auth via Authorization header OR ?token= query param (needed + because cannot send custom headers). """ from fastapi.responses import FileResponse @@ -3412,6 +4577,10 @@ async def serve_sandbox_file( safe_filename = os.path.basename(filename) if not safe_filename or safe_filename in (".", ".."): raise HTTPException(status_code = 404, detail = "Not found") + # Defense-in-depth allowlist (clears CodeQL py/path-injection), still allowing + # names like "loss curve.png"; basename + extension + realpath below are the guards. + if not _re.fullmatch(r"[^/\\\x00-\x1f]{1,255}", safe_filename): + raise HTTPException(status_code = 404, detail = "Not found") # ── Extension allowlist ───────────────────────────────────── ext = os.path.splitext(safe_filename)[1].lower() @@ -3423,16 +4592,11 @@ async def serve_sandbox_file( ) # ── Path containment check ────────────────────────────────── - home = os.path.expanduser("~") - sandbox_root = os.path.realpath(os.path.join(home, "studio_sandbox")) - safe_session = os.path.basename(session_id.replace("..", "")) - if not safe_session: - raise HTTPException(status_code = 404, detail = "Not found") + from core.inference.tools import get_sandbox_workdir - file_path = os.path.realpath( - os.path.join(sandbox_root, safe_session, safe_filename) - ) - if not file_path.startswith(sandbox_root + os.sep): + sandbox_dir = os.path.realpath(get_sandbox_workdir(session_id)) + file_path = os.path.realpath(os.path.join(sandbox_dir, safe_filename)) + if file_path != sandbox_dir and not file_path.startswith(sandbox_dir + os.sep): raise HTTPException( status_code = status.HTTP_403_FORBIDDEN, detail = "Access denied", @@ -3456,41 +4620,93 @@ async def serve_sandbox_file( # ===================================================================== +def _openai_model_objects() -> list[dict]: + """The model objects GET /v1/models exposes (one per loaded local backend). + + Shared by the LIST and RETRIEVE handlers so both report the same ids and + field shape. + """ + models: list[dict] = [] + _created = int(time.time()) + + # Check GGUF backend + llama_backend = get_llama_cpp_backend() + if llama_backend.is_loaded: + entry = { + "id": llama_backend.model_identifier, + "object": "model", + "created": _created, + "owned_by": "local", + } + _ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None)) + if _ctx is not None: + entry["context_length"] = _ctx + _max_ctx = _positive_int_or_none(getattr(llama_backend, "max_context_length", None)) + if _max_ctx is not None: + entry["max_context_length"] = _max_ctx + _native_ctx = _positive_int_or_none(getattr(llama_backend, "native_context_length", None)) + if _native_ctx is not None: + entry["native_context_length"] = _native_ctx + models.append(entry) + + # Check Unsloth backend + backend = get_inference_backend() + if backend.active_model_name: + model_info = backend.models.get(backend.active_model_name, {}) + entry = { + "id": backend.active_model_name, + "object": "model", + "created": _created, + "owned_by": "local", + } + _ctx = _positive_int_or_none(model_info.get("context_length")) + if _ctx is None: + for _candidate in ( + getattr(backend, "context_length", None), + getattr(backend, "max_seq_length", None), + ): + _ctx = _positive_int_or_none(_candidate) + if _ctx is not None: + break + if _ctx is not None: + entry["context_length"] = _ctx + models.append(entry) + + return models + + @router.get("/models") -async def openai_list_models( - current_subject: str = Depends(get_current_subject), -): +async def openai_list_models(current_subject: str = Depends(get_current_subject)): """ OpenAI-compatible model listing endpoint. Returns the currently loaded model in the format expected by OpenAI-compatible clients (``GET /v1/models``). """ - models = [] + return {"object": "list", "data": _openai_model_objects()} - # Check GGUF backend - llama_backend = get_llama_cpp_backend() - if llama_backend.is_loaded: - models.append( - { - "id": llama_backend.model_identifier, - "object": "model", - "owned_by": "local", - } - ) - # Check Unsloth backend - backend = get_inference_backend() - if backend.active_model_name: - models.append( - { - "id": backend.active_model_name, - "object": "model", - "owned_by": "local", - } - ) +@router.get("/models/{model_id:path}") +async def openai_retrieve_model(model_id: str, current_subject: str = Depends(get_current_subject)): + """ + OpenAI-compatible single-model retrieval endpoint (``GET /v1/models/{id}``). - return {"object": "list", "data": models} + Returns the bare model object when ``model_id`` matches a loaded local + model, or 404 model_not_found otherwise. Defined after the LIST route so + it does not shadow it; ``{model_id:path}`` keeps ids with slashes intact. + """ + for model in _openai_model_objects(): + if model["id"] == model_id: + return model + raise HTTPException( + status_code = 404, + detail = openai_error_body( + f"The model '{model_id}' does not exist", + status = 404, + code = "model_not_found", + param = "id", + ), + ) # ===================================================================== @@ -3499,15 +4715,12 @@ async def openai_list_models( @router.post("/completions") -async def openai_completions( - request: Request, - current_subject: str = Depends(get_current_subject), -): +async def openai_completions(request: Request, current_subject: str = Depends(get_current_subject)): """ OpenAI-compatible text completions endpoint (non-chat). - Transparently proxies to the running llama-server's ``/v1/completions``. - Only available when a GGUF model is loaded. + Proxies to the running llama-server's ``/v1/completions``. Only available + when a GGUF model is loaded. """ llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: @@ -3523,15 +4736,17 @@ async def openai_completions( if is_stream: async def _stream(): - # Manual httpx client/response lifecycle AND explicit - # aiter_bytes() iterator close — see _anthropic_passthrough_stream - # for the full rationale. Saving `bytes_iter = resp.aiter_bytes()` - # and `await bytes_iter.aclose()` in the finally block is the - # part that matters for avoiding the Python 3.13 + httpcore - # 1.0.x "Exception ignored in: " / anyio - # cancel-scope trace: an anonymous async for leaves the - # iterator unclosed, so Python's asyncgen GC finalizer runs - # cleanup on a later pass in a different asyncio task. + # Manual httpx client/response lifecycle AND explicit iterator + # close — see _anthropic_passthrough_stream for the full rationale. + # Saving the iterator and closing it in the finally block avoids the + # Python 3.13 + httpcore 1.0.x "Exception ignored in: + # " / anyio cancel-scope trace. + # + # Buffer the relay into whole SSE events (split on the blank-line + # separator) so _cmpl_stream_event_out can rewrite the cmpl- id and + # honor stream_options.include_usage per event, while keeping SSE + # framing and token bytes intact. + _include_usage = bool((body.get("stream_options") or {}).get("include_usage")) client = httpx.AsyncClient(timeout = 600) resp = None bytes_iter = None @@ -3539,8 +4754,21 @@ async def openai_completions( req = client.build_request("POST", target_url, json = body) resp = await client.send(req, stream = True) bytes_iter = resp.aiter_bytes() + buffer = b"" async for chunk in bytes_iter: - yield chunk + buffer += chunk + while b"\n\n" in buffer: + event, buffer = buffer.split(b"\n\n", 1) + out = _cmpl_stream_event_out(event, _include_usage) + if out is not None: + yield out + b"\n\n" + if buffer: + out = _cmpl_stream_event_out(buffer, _include_usage) + if out is not None: + # Re-add the SSE separator the split consumed, so a final + # event arriving without a trailing blank line is still + # terminated for the client's parser. + yield out + b"\n\n" except Exception as e: logger.error("openai_completions stream error: %s", e) finally: @@ -3563,8 +4791,12 @@ async def openai_completions( else: async with httpx.AsyncClient() as client: resp = await client.post(target_url, json = body, timeout = 600) + + if resp.status_code != 200: + raise _openai_passthrough_error(resp.status_code, resp.text) + return Response( - content = resp.content, + content = _rewrite_cmpl_id(resp.content), status_code = resp.status_code, media_type = "application/json", ) @@ -3576,17 +4808,14 @@ async def openai_completions( @router.post("/embeddings") -async def openai_embeddings( - request: Request, - current_subject: str = Depends(get_current_subject), -): +async def openai_embeddings(request: Request, current_subject: str = Depends(get_current_subject)): """ OpenAI-compatible embeddings endpoint. - Transparently proxies to the running llama-server's ``/v1/embeddings``. - Only available when a GGUF model is loaded. - Note: the loaded model must support pooling; otherwise llama-server - will return an error (expected). + Proxies to the running llama-server's ``/v1/embeddings``. Only available + when a GGUF model is loaded. + Note: the loaded model must support pooling, else llama-server returns an + error (expected). """ llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: @@ -3612,9 +4841,7 @@ async def openai_embeddings( # ===================================================================== -def _translate_responses_tools_to_chat( - tools: Optional[list[dict]], -) -> Optional[list[dict]]: +def _translate_responses_tools_to_chat(tools: Optional[list[dict]]) -> Optional[list[dict]]: """Translate Responses-shape function tools to the Chat Completions nested shape. Responses uses a flat shape per tool entry:: @@ -3629,9 +4856,9 @@ def _translate_responses_tools_to_chat( "parameters": {...}, "strict": true}} Only ``type=="function"`` entries are forwarded. Built-in Responses tools - (``web_search``, ``file_search``, ``mcp``, ...) are dropped because - llama-server does not implement them server-side; keeping them in the - request would produce an opaque upstream 400. + (``web_search``, ``file_search``, ``mcp``, ...) are dropped: llama-server + doesn't implement them server-side, so keeping them would produce an opaque + upstream 400. """ if not tools: return None @@ -3658,8 +4885,8 @@ def _translate_responses_tool_choice_to_chat(tool_choice: Any) -> Any: """Translate a Responses-shape ``tool_choice`` to the Chat Completions shape. String values (``"auto"``/``"none"``/``"required"``) pass through unchanged. - The Responses forcing object ``{"type": "function", "name": "X"}`` is - converted to Chat Completions' ``{"type": "function", "function": {"name": "X"}}``. + The Responses forcing object ``{"type": "function", "name": "X"}`` becomes + Chat Completions' ``{"type": "function", "function": {"name": "X"}}``. Unknown / built-in tool choices are forwarded as-is; llama-server ignores what it doesn't recognise. """ @@ -3693,29 +4920,38 @@ def _responses_message_text(content: Union[str, list]) -> str: return "\n".join(parts) +def _responses_tool_output_text(output: Union[str, list]) -> str: + """Return Chat Completions-safe content for a Responses tool result.""" + if isinstance(output, str): + return output if output.strip() else "(no output)" + + if output: + return json.dumps(output) + + return "(no output)" + + def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: - """Convert a ResponsesRequest's ``input`` into Chat-format ``ChatMessage`` list. + """Convert a ResponsesRequest's ``input`` into a Chat-format ``ChatMessage`` list. Handles the three input item shapes allowed by the Responses API: - - ``ResponsesInputMessage`` — regular chat messages (text or multimodal). - - ``ResponsesFunctionCallInputItem`` — a prior assistant tool call replayed - on a follow-up turn. Converted into an assistant message carrying a + - ``ResponsesInputMessage`` -- regular chat messages (text or multimodal). + - ``ResponsesFunctionCallInputItem`` -- a prior assistant tool call + replayed on a follow-up turn. Becomes an assistant message carrying a Chat Completions ``tool_calls`` entry keyed by ``call_id``. - - ``ResponsesFunctionCallOutputInputItem`` — a tool result the client is - returning. Converted into a ``role="tool"`` message with ``tool_call_id`` - set to the originating ``call_id`` so llama-server can reconcile the - call with its result. + - ``ResponsesFunctionCallOutputInputItem`` -- a tool result the client is + returning. Becomes a ``role="tool"`` message with ``tool_call_id`` set to + the originating ``call_id`` so llama-server can reconcile call with result. - System / developer content is collected from ``instructions`` *and* from - any ``role="system"`` / ``role="developer"`` entries in ``input``, then - merged into a single ``role="system"`` message placed at the top of the - returned list. This satisfies strict chat templates (harmony / gpt-oss, - Qwen3, ...) whose Jinja raises ``"System message must be at the - beginning."`` when more than one system message is present or when a - system message appears after a user turn — the exact pattern the OpenAI - Codex CLI hits, since Codex sets ``instructions`` *and* also sends a - developer message in ``input``. + System / developer content is collected from ``instructions`` *and* any + ``role="system"`` / ``role="developer"`` entries in ``input``, then merged + into a single top-of-list ``role="system"`` message. This satisfies strict + chat templates (harmony / gpt-oss, Qwen3, ...) whose Jinja raises + ``"System message must be at the beginning."`` when more than one system + message is present or a system message follows a user turn -- the exact + pattern the OpenAI Codex CLI hits, since Codex sets ``instructions`` *and* + also sends a developer message in ``input``. """ system_parts: list[str] = [] messages: list[ChatMessage] = [] @@ -3753,11 +4989,10 @@ def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: continue if isinstance(item, ResponsesFunctionCallOutputInputItem): - # Chat Completions `role="tool"` requires a string content; if a - # Responses client sends a content-array output, serialize it. - output = item.output - if not isinstance(output, str): - output = json.dumps(output) + # Chat Completions `role="tool"` requires string content; serialize + # a Responses content-array output and keep empty outputs from + # tripping the stricter ChatMessage role validator. + output = _responses_tool_output_text(item.output) messages.append( ChatMessage( role = "tool", @@ -3768,13 +5003,13 @@ def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: continue if isinstance(item, ResponsesUnknownInputItem): - # Reasoning items and any other unmodelled top-level Responses - # item types are silently dropped — llama-server-backed GGUFs - # cannot consume them and our lenient validation let them in so - # unrelated turns don't 422. + # Reasoning items and other unmodelled top-level Responses item + # types are silently dropped -- llama-server-backed GGUFs can't + # consume them; lenient validation lets them in so unrelated turns + # don't 422. continue - # ResponsesInputMessage — hoist system/developer to the top, merge. + # ResponsesInputMessage -- hoist system/developer to the top, merge. if item.role in ("system", "developer"): hoisted = _responses_message_text(item.content) if hoisted: @@ -3787,16 +5022,16 @@ def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: # Assistant-replay turns come back as content = [output_text, ...]. # Chat Completions' assistant role expects a plain string, not a - # multimodal content array, so flatten output_text (and any stray - # input_text / unknown text) to a single string. + # multimodal array, so flatten output_text (and any stray input_text / + # unknown text) to a single string. if item.role == "assistant": text = _responses_message_text(item.content) if text: messages.append(ChatMessage(role = "assistant", content = text)) continue - # User (and any other remaining roles) — keep multimodal when - # present, drop unknown content parts silently. + # User (and any other remaining roles) -- keep multimodal when present, + # drop unknown content parts silently. parts: list = [] for part in item.content: if isinstance(part, (ResponsesInputTextPart, ResponsesOutputTextPart)): @@ -3810,9 +5045,8 @@ def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: ) # ResponsesUnknownContentPart and anything else: drop. if parts: - # Collapse single-text-part content to a plain string so roles - # that reject multimodal arrays (e.g. legacy templates) still - # accept the message. + # Collapse single-text-part content to a plain string so roles that + # reject multimodal arrays (e.g. legacy templates) still accept it. if len(parts) == 1 and isinstance(parts[0], TextContentPart): messages.append(ChatMessage(role = item.role, content = parts[0].text)) else: @@ -3831,8 +5065,7 @@ def _build_chat_request( Tools and ``tool_choice`` are translated from the flat Responses shape to the nested Chat Completions shape here so the existing #5099 - ``/v1/chat/completions`` client-side pass-through picks them up without - further modification. + ``/v1/chat/completions`` client-side pass-through picks them up unchanged. """ chat_kwargs: dict = dict( model = payload.model, @@ -3853,23 +5086,18 @@ def _build_chat_request( chat_tool_choice = _translate_responses_tool_choice_to_chat(payload.tool_choice) if chat_tool_choice is not None: chat_kwargs["tool_choice"] = chat_tool_choice + if payload.parallel_tool_calls is not None: + chat_kwargs["parallel_tool_calls"] = payload.parallel_tool_calls - req = ChatCompletionRequest(**chat_kwargs) - # `parallel_tool_calls` is not a first-class field on ChatCompletionRequest, - # but the model allows extras and _build_openai_passthrough_body forwards - # only explicitly-known fields. Llama-server does not currently implement - # parallel_tool_calls semantics, so we accept-and-ignore it on the - # Responses side to avoid breaking SDK clients that always send it. - return req + return ChatCompletionRequest(**chat_kwargs) def _chat_tool_calls_to_responses_output(tool_calls: list[dict]) -> list[dict]: """Map Chat Completions ``tool_calls`` into Responses ``function_call`` output items. The Chat Completions id (``call_xxx``) is the shared correlation key across - turns in the OpenAI Responses API — it is stored as ``call_id`` on the - output item and must be echoed back by the client as - ``function_call_output.call_id`` on the next turn. + turns in the Responses API -- stored as ``call_id`` on the output item and + echoed back by the client as ``function_call_output.call_id`` next turn. """ items: list[dict] = [] for tc in tool_calls: @@ -3888,15 +5116,13 @@ def _chat_tool_calls_to_responses_output(tool_calls: list[dict]) -> list[dict]: async def _responses_non_streaming( - payload: ResponsesRequest, - messages: list[ChatMessage], - request: Request, + payload: ResponsesRequest, messages: list[ChatMessage], request: Request ) -> JSONResponse: """Handle a non-streaming Responses API call.""" chat_req = _build_chat_request(payload, messages, stream = False) result = await openai_chat_completions(chat_req, request) - # openai_chat_completions returns a JSONResponse for non-streaming + # openai_chat_completions returns a JSONResponse for non-streaming. if isinstance(result, JSONResponse): body = json.loads(result.body.decode()) elif isinstance(result, Response): @@ -3919,10 +5145,9 @@ async def _responses_non_streaming( resp_id = f"resp_{uuid.uuid4().hex[:12]}" # Responses API emits each tool call as its own top-level output item, - # alongside an optional assistant text message. Emit the text message - # only when the model actually produced content, so clients that expect - # a pure tool-call turn (finish_reason="tool_calls") don't see a spurious - # empty message item. + # plus an optional assistant text message. Emit the text message only when + # the model produced content, so clients expecting a pure tool-call turn + # (finish_reason="tool_calls") don't see a spurious empty message item. output_items: list[dict] = [] if text: msg_id = f"msg_{uuid.uuid4().hex[:12]}" @@ -3956,31 +5181,28 @@ async def _responses_non_streaming( async def _responses_stream( - payload: ResponsesRequest, - messages: list[ChatMessage], - request: Request, + payload: ResponsesRequest, messages: list[ChatMessage], request: Request ): """Handle a streaming Responses API call, emitting named SSE events. For GGUF models the request goes directly to llama-server's - ``/v1/chat/completions`` endpoint from inside the StreamingResponse - child task — a single httpx lifecycle, a single async generator. - Wrapping the existing ``openai_chat_completions`` pass-through (which - already does its own httpx lifecycle) stacks two generators: Python - 3.13 + httpcore 1.0.x then loses the close-propagation chain on the - innermost ``HTTP11ConnectionByteStream`` at asyncgen finalisation, - tripping "Attempted to exit cancel scope in a different task" / - "async generator ignored GeneratorExit". The direct path avoids that - altogether. Non-GGUF falls back to the wrapper (which doesn't use - httpx, so the issue doesn't apply). + ``/v1/chat/completions`` from inside the StreamingResponse child task -- one + httpx lifecycle, one async generator. Wrapping the existing + ``openai_chat_completions`` pass-through (which has its own httpx lifecycle) + stacks two generators: Python 3.13 + httpcore 1.0.x then loses the + close-propagation chain on the innermost ``HTTP11ConnectionByteStream`` at + asyncgen finalisation, tripping "Attempted to exit cancel scope in a + different task" / "async generator ignored GeneratorExit". The direct path + avoids that. Non-GGUF falls back to the wrapper (which doesn't use httpx, so + the issue doesn't apply). Text deltas arrive as ``response.output_text.delta`` on a single ``message`` output item at ``output_index=0``. Each tool call from ``delta.tool_calls[]`` is promoted to its own top-level ``function_call`` - output item (one per distinct ``tool_calls[].index``), and relayed as + output item (one per distinct ``tool_calls[].index``) and relayed as ``response.function_call_arguments.delta`` / ``.done`` events so clients - (Codex, OpenAI Python SDK) can reconstruct the call incrementally and - reply with a ``function_call_output`` item on the next turn. + (Codex, OpenAI Python SDK) can reconstruct the call incrementally and reply + with a ``function_call_output`` item next turn. """ resp_id = f"resp_{uuid.uuid4().hex[:12]}" msg_id = f"msg_{uuid.uuid4().hex[:12]}" @@ -3990,13 +5212,12 @@ async def _responses_stream( llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: - # The direct pass-through is GGUF-only. Non-GGUF /v1/responses - # streaming isn't a Codex-compatible path today and wrapping the - # transformers backend's streaming generator here would re- - # introduce the double-layer asyncgen close pattern that produces - # "Attempted to exit cancel scope in a different task" on Python - # 3.13. Surface a typed 400 so the client sees a useful error - # instead of a dangling stream. + # The direct pass-through is GGUF-only. Non-GGUF /v1/responses streaming + # isn't a Codex-compatible path today, and wrapping the transformers + # backend's streaming generator here would re-introduce the + # double-layer asyncgen close pattern that produces "Attempted to exit + # cancel scope in a different task" on Python 3.13. Surface a typed 400 + # so the client sees a useful error instead of a dangling stream. raise HTTPException( status_code = 400, detail = ( @@ -4008,8 +5229,7 @@ async def _responses_stream( # Direct pass-through bypasses the openai_chat_completions image gate. if not llama_backend.is_vision and any( - isinstance(m.content, list) - and any(isinstance(p, ImageContentPart) for p in m.content) + isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content) for m in messages ): raise HTTPException( @@ -4018,16 +5238,17 @@ async def _responses_stream( ) body = _build_openai_passthrough_body( - chat_req, backend_ctx = llama_backend.context_length + chat_req, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) + body["stream_options"] = {"include_usage": True} target_url = f"{llama_backend.base_url}/v1/chat/completions" async def event_generator(): full_text = "" input_tokens = 0 output_tokens = 0 - # Per-tool-call state keyed by the Chat Completions `tool_calls[].index` - # which stays stable across chunks for the same call. Values are: + # Per-tool-call state keyed by Chat Completions `tool_calls[].index`, + # stable across chunks for the same call. Values: # {output_index, item_id, call_id, name, arguments, opened} tool_call_state: dict[int, dict] = {} # Text message lives at output_index 0; tool calls claim 1, 2, ... @@ -4081,11 +5302,11 @@ async def _responses_stream( yield f"event: response.content_part.added\ndata: {json.dumps({'type': 'response.content_part.added', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': content_part})}\n\n" # ── Direct httpx lifecycle to llama-server ── - # Full same-task open + close, identical pattern to - # _openai_passthrough_stream and _anthropic_passthrough_stream: - # no `async with`, explicit aclose of lines_iter BEFORE resp / - # client so the innermost httpcore byte stream is finalised in - # this task (not via Python's asyncgen GC in a sibling task). + # Full same-task open + close, same pattern as + # _openai_passthrough_stream and _anthropic_passthrough_stream: no + # `async with`, explicit aclose of lines_iter BEFORE resp / client so + # the innermost httpcore byte stream is finalised in this task (not via + # the asyncgen GC in a sibling task). client = httpx.AsyncClient(timeout = 600) resp = None lines_iter = None @@ -4124,6 +5345,8 @@ async def _responses_stream( chunk_data = json.loads(data_str) except json.JSONDecodeError: continue + if payload.parallel_tool_calls is False: + _drop_parallel_tool_call_deltas(chunk_data) choices = chunk_data.get("choices", []) if not choices: @@ -4151,7 +5374,7 @@ async def _responses_stream( st = tool_call_state.get(idx) fn = tc.get("function") or {} if st is None: - # First chunk for this tool call — allocate an + # First chunk for this tool call -- allocate an # output_index and emit output_item.added. st = { "output_index": next_output_index, @@ -4164,8 +5387,8 @@ async def _responses_stream( next_output_index += 1 tool_call_state[idx] = st else: - # Later chunks sometimes carry the id/name only - # once; merge when present. + # Later chunks sometimes carry id/name only once; merge + # when present. if tc.get("id") and not st["call_id"]: st["call_id"] = tc["id"] if fn.get("name") and not st["name"]: @@ -4198,9 +5421,9 @@ async def _responses_stream( } yield f"event: response.function_call_arguments.delta\ndata: {json.dumps(args_delta_event)}\n\n" elif arg_delta: - # Buffer the args until we can open the item - # (id/name arrive in the same chunk as the first - # arg delta for some models — but if not, stash). + # Buffer args until we can open the item (some models + # send id/name in the same chunk as the first arg delta; + # if not, stash). st["arguments"] += arg_delta usage = chunk_data.get("usage") @@ -4227,8 +5450,8 @@ async def _responses_stream( # ── Closing events for tool calls ── for st in sorted(tool_call_state.values(), key = lambda s: s["output_index"]): - # If id/name never arrived (malformed upstream), synthesise so - # the client still sees a coherent frame sequence. + # If id/name never arrived (malformed upstream), synthesise so the + # client still sees a coherent frame sequence. if not st["opened"]: if not st["call_id"]: st["call_id"] = f"call_{uuid.uuid4().hex[:12]}" @@ -4331,10 +5554,9 @@ async def openai_responses( """ OpenAI Responses API endpoint. - Accepts the Responses-format request, converts it to a - ChatCompletionRequest internally, and returns a response - matching the OpenAI Responses API schema (output array, - input_tokens/output_tokens, named SSE events for streaming). + Accepts a Responses-format request, converts it to a ChatCompletionRequest + internally, and returns a response matching the Responses API schema + (output array, input_tokens/output_tokens, named SSE events for streaming). """ messages = _normalise_responses_input(payload) if not messages: @@ -4368,9 +5590,9 @@ def _anthropic_requested_studio_tools(tools: Optional[list]) -> set[str]: # Client tools always carry input_schema; server tools never do. if td.get("input_schema") is not None: continue - # Anthropic dispatches server tools by `type` (not by bare `name`); - # matching name too would let a malformed client tool like - # `{"name": "python"}` silently flip into server-execution mode. + # Anthropic dispatches server tools by `type`, not bare `name`; matching + # name too would let a malformed client tool like `{"name": "python"}` + # silently flip into server-execution mode. type_ = td.get("type") if isinstance(type_, str) and type_ in _STUDIO_ANTHROPIC_TOOL_ALIASES: requested.add(_STUDIO_ANTHROPIC_TOOL_ALIASES[type_]) @@ -4378,9 +5600,7 @@ def _anthropic_requested_studio_tools(tools: Optional[list]) -> set[str]: def _select_anthropic_server_tools( - all_tools: list[dict], - requested_studio_tools: set[str], - enabled_tools: Optional[list[str]], + all_tools: list[dict], requested_studio_tools: set[str], enabled_tools: Optional[list[str]] ) -> list[dict]: """Select Studio tools requested through Anthropic tools and extensions.""" if not requested_studio_tools and enabled_tools is None: @@ -4393,22 +5613,19 @@ def _select_anthropic_server_tools( return [tool for tool in all_tools if tool["function"]["name"] in selected_names] -def _normalize_anthropic_openai_images( - openai_messages: list[dict], is_vision: bool -) -> bool: - """Enforce the vision guard on translated Anthropic messages and - normalize any ``image_url`` parts with base64 data URLs to PNG. +def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: bool) -> bool: + """Enforce the vision guard on translated Anthropic messages and normalize + any base64-data-URL ``image_url`` parts to PNG. llama-server's stb_image only handles a few formats (JPEG/PNG/BMP/…); - Anthropic clients commonly send JPEG or WebP, and Claude Code sends - WebP. Re-encoding everything to PNG mirrors the behavior of - `_openai_messages_for_passthrough` / the GGUF branch of - `/v1/chat/completions` so the two endpoints agree. + Anthropic clients commonly send JPEG or WebP, and Claude Code sends WebP. + Re-encoding everything to PNG mirrors `_openai_messages_for_passthrough` / + the GGUF branch of `/v1/chat/completions` so the two endpoints agree. - Mutates ``openai_messages`` in place. Returns ``True`` when any - image part was seen (so the caller can skip a second scan). Raises - HTTPException(400) when images are present but the active model is - not a vision model, or when an image cannot be decoded. + Mutates ``openai_messages`` in place. Returns ``True`` when any image part + was seen (so the caller can skip a second scan). Raises HTTPException(400) + when images are present but the active model isn't a vision model, or when + an image cannot be decoded. """ from PIL import Image @@ -4451,6 +5668,70 @@ def _normalize_anthropic_openai_images( return has_image +@router.post("/messages/count_tokens") +async def anthropic_count_tokens( + payload: AnthropicMessagesRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """Anthropic-compatible token-counting endpoint (POST /v1/messages/count_tokens). + + Translates the Anthropic request to OpenAI form (the same translation the + /messages handler uses), counts prompt tokens with the loaded GGUF model's + tokenizer, and returns ``{"input_tokens": int}`` only. Unlike /messages, + max_tokens is NOT required here. + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + # Same Anthropic → OpenAI translation as anthropic_messages: system is + # folded into the messages list, so pass system=None to the counter. + openai_messages = anthropic_messages_to_openai( + [m.model_dump() for m in payload.messages], + payload.system, + ) + # Apply the same sanitization /messages does before generation, so the count + # matches the prompt the real request would build (otherwise empty-assistant + # sentinels / synthetic tool history inflate the count or hit the fallback). + openai_messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels(openai_messages) + ) + openai_tools = anthropic_tools_to_openai(payload.tools or []) or None + + try: + count = await asyncio.to_thread( + llama_backend.count_chat_tokens, + openai_messages, + None, + openai_tools, + strict = True, + ) + except Exception: + raise HTTPException( + status_code = 503, + detail = "Unable to count tokens with the loaded model tokenizer.", + ) + return JSONResponse(content = {"input_tokens": int(count)}) + + +def _set_or_prepend_system_message( + messages: Optional[list[dict]], system_prompt: str +) -> list[dict]: + """Return messages with a single leading system prompt, preserving multimodal parts.""" + safe_messages = messages or [] + if not system_prompt: + return safe_messages + + # Drop existing system/developer turns so the backend never sees duplicate + # or conflicting system instructions, then prepend the resolved prompt. + others = [dict(msg) for msg in safe_messages if msg.get("role") not in ("system", "developer")] + return [{"role": "system", "content": system_prompt}, *others] + + @router.post("/messages") async def anthropic_messages( payload: AnthropicMessagesRequest, @@ -4460,10 +5741,10 @@ async def anthropic_messages( """ Anthropic-compatible Messages API endpoint. - Translates Anthropic message format to internal OpenAI format, runs - through the existing agentic tool loop when tools are provided, and - returns responses in Anthropic Messages API format (streaming SSE or - non-streaming JSON). + Translates Anthropic message format to internal OpenAI format, runs through + the existing agentic tool loop when tools are provided, and returns + responses in Anthropic Messages API format (streaming SSE or non-streaming + JSON). """ llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: @@ -4472,6 +5753,18 @@ async def anthropic_messages( detail = "No GGUF model loaded. Load a GGUF model first.", ) + # max_tokens is a required field on the Anthropic Messages API; real + # Anthropic returns a 400 invalid_request_error when it is omitted. + if payload.max_tokens is None: + raise HTTPException( + status_code = 400, + detail = anthropic_error_body( + "max_tokens: field required", + status = 400, + err_type = "invalid_request_error", + ), + ) + model_name = getattr(llama_backend, "model_identifier", None) or payload.model message_id = f"msg_{uuid.uuid4().hex[:24]}" @@ -4480,14 +5773,21 @@ async def anthropic_messages( [m.model_dump() for m in payload.messages], payload.system, ) - openai_messages = _drop_empty_assistant_sentinels(openai_messages) - - # Enforce vision guard + re-encode embedded images to PNG so the - # Anthropic endpoint matches the behavior of /v1/chat/completions. - _has_image = _normalize_anthropic_openai_images( - openai_messages, llama_backend.is_vision + # Strip synthetic provider-side builtin tool history (web_search, + # web_fetch, code_execution, image_generation cards tagged with + # _server_tool or extra_content.google.native_part) before handing off to + # local llama-server. The local /v1/chat/completions and GGUF passthrough + # builders apply the same strip; without it an Anthropic /v1/messages caller + # replaying a prior provider-side tool_use forwards fake builtin tool + # history to a backend with no matching function declarations. + openai_messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels(openai_messages) ) + # Enforce vision guard + re-encode embedded images to PNG so the Anthropic + # endpoint matches /v1/chat/completions. + _has_image = _normalize_anthropic_openai_images(openai_messages, llama_backend.is_vision) + temperature = payload.temperature if payload.temperature is not None else 0.6 top_p = payload.top_p if payload.top_p is not None else 0.95 top_k = payload.top_k if payload.top_k is not None else 20 @@ -4495,14 +5795,11 @@ async def anthropic_messages( repetition_penalty = ( payload.repetition_penalty if payload.repetition_penalty is not None else 1.0 ) - presence_penalty = ( - payload.presence_penalty if payload.presence_penalty is not None else 0.0 - ) + presence_penalty = payload.presence_penalty if payload.presence_penalty is not None else 0.0 stop = payload.stop_sequences or None - # Translate Anthropic tool_choice to OpenAI format for forwarding to - # llama-server. Falls back to "auto" when unset or unrecognized, which - # matches the prior hardcoded behavior. + # Translate Anthropic tool_choice to OpenAI format for llama-server. Falls + # back to "auto" when unset or unrecognized (prior hardcoded behavior). openai_tool_choice = anthropic_tool_choice_to_openai(payload.tool_choice) if openai_tool_choice is None: openai_tool_choice = "auto" @@ -4514,17 +5811,17 @@ async def anthropic_messages( # 1. enable_tools=true → server-side execution of built-in tools (Unsloth shorthand) # 2. tools=[...] only → client-side pass-through (standard Anthropic behavior) # 3. neither → plain chat - # Server-side agentic loop doesn't support multimodal input — matches + # The server-side agentic loop doesn't support multimodal input -- matches # the `not image_b64` gate in /v1/chat/completions. requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) - # Reject malformed client tools at the boundary. AnthropicTool was - # relaxed to Optional[name]/Optional[input_schema] for server tools, - # so the converter silently drops incomplete entries — surface them - # as 400. A `type` field marks a server-tool declaration per spec - # (unrecognized server tools are accepted as no-ops); anything else - # without input_schema or name is malformed and must not be allowed - # to silently flip execution mode or disable tool calling. + # Reject malformed client tools at the boundary. AnthropicTool was relaxed + # to Optional[name]/Optional[input_schema] for server tools, so the + # converter silently drops incomplete entries -- surface them as 400. A + # `type` field marks a server-tool declaration per spec (unrecognized server + # tools are accepted as no-ops); anything else without input_schema or name + # is malformed and must not be allowed to silently flip execution mode or + # disable tool calling. for tool in payload.tools or []: td = tool if isinstance(tool, dict) else tool.model_dump() name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") @@ -4539,17 +5836,17 @@ async def anthropic_messages( detail = "Client tool is missing required field 'name'.", ) - # Detect client tools from the raw payload (presence of input_schema) - # so the mixed-mode check below isn't fooled by a name collision with - # a server-tool alias that the post-filter would silently drop. + # Detect client tools from the raw payload (presence of input_schema) so the + # mixed-mode check below isn't fooled by a name collision with a server-tool + # alias that the post-filter would silently drop. _has_client_tool = any( (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None for t in payload.tools or [] ) - # The server-tool agentic loop executes tools in-process and cannot - # relay unknown client functions back to the caller, so mixed requests - # would silently drop the client tools. Reject explicitly instead. + # The server-tool agentic loop executes tools in-process and can't relay + # unknown client functions back to the caller, so mixed requests would + # silently drop the client tools. Reject explicitly instead. if requested_studio_tools and _has_client_tool: raise HTTPException( status_code = 400, @@ -4566,9 +5863,9 @@ async def anthropic_messages( if tool.get("function", {}).get("name") not in requested_studio_tools ] - # An Anthropic server-tool declaration implies server-tool mode, but - # only when tools aren't explicitly disabled (CLI --disable-tools or - # per-request enable_tools=false). Explicit False always wins. + # An Anthropic server-tool declaration implies server-tool mode, but only + # when tools aren't explicitly disabled (CLI --disable-tools or per-request + # enable_tools=false). Explicit False always wins. _enable = _effective_enable_tools(payload) server_tools = ( (_enable or (_enable is None and bool(requested_studio_tools))) @@ -4576,9 +5873,15 @@ async def anthropic_messages( and not _has_image ) client_tools = ( - not server_tools - and len(openai_client_tools) > 0 - and llama_backend.supports_tools + not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools + ) + + # Anthropic tool_choice.disable_parallel_tool_use caps the response to a + # single tool_use block. Computed here so BOTH the client-tool passthrough + # and the server-tool path honor it. + _disable_parallel = bool( + isinstance(payload.tool_choice, dict) + and payload.tool_choice.get("disable_parallel_tool_use") ) # ── Client-side pass-through path ───────────────────────── @@ -4605,6 +5908,7 @@ async def anthropic_messages( tool_choice = openai_tool_choice, session_id = payload.session_id, cancel_id = payload.cancel_id, + disable_parallel_tool_use = _disable_parallel, ) return await _anthropic_passthrough_non_streaming( llama_backend, @@ -4621,6 +5925,7 @@ async def anthropic_messages( repetition_penalty = repetition_penalty, presence_penalty = presence_penalty, tool_choice = openai_tool_choice, + disable_parallel_tool_use = _disable_parallel, ) if server_tools: @@ -4633,55 +5938,12 @@ async def anthropic_messages( ) # Build tool-use system prompt nudge (same logic as /chat/completions) - _tool_names = {t["function"]["name"] for t in openai_tools} - _has_web = "web_search" in _tool_names - _has_code = "python" in _tool_names or "terminal" in _tool_names - - _date_line = f"The current date is {_date.today().isoformat()}." - _model_size_b = _extract_model_size_b(model_name) - _is_small_model = _model_size_b is not None and _model_size_b < 9 - - if _is_small_model: - _web_tips = "Do not repeat the same search query." - else: - _web_tips = ( - "When you search and find a relevant URL in the results, " - "fetch its full content by calling web_search with the url parameter. " - "Do not repeat the same search query. If a search returns " - "no useful results, try rephrasing or fetching a result URL directly." - ) - _code_tips = ( - "Use code execution for math, calculations, data processing, " - "or to parse and analyze information from tool results." + _nudge = _build_tool_action_nudge( + tools = openai_tools, + model_name = model_name, ) - if _has_web and _has_code: - _nudge = ( - _date_line + " " - "You have access to tools. When appropriate, prefer using " - "tools rather than answering from memory. " - + _web_tips - + " " - + _code_tips - ) - elif _has_code: - _nudge = ( - _date_line + " " - "You have access to tools. When appropriate, prefer using " - "code execution rather than answering from memory. " + _code_tips - ) - elif _has_web: - _nudge = ( - _date_line + " " - "You have access to tools. When appropriate, prefer using " - "web search for up-to-date or uncertain factual " - "information rather than answering from memory. " + _web_tips - ) - else: - _nudge = "" - if _nudge: - _nudge += _TOOL_ACTION_NUDGE # Inject into system prompt if openai_messages and openai_messages[0].get("role") == "system": openai_messages[0]["content"] = ( @@ -4712,6 +5974,9 @@ async def anthropic_messages( auto_heal_tool_calls = True, tool_call_timeout = 300, session_id = payload.session_id, + # Anthropic passthrough has no rag_scope field (RAG is local-only). + rag_scope = getattr(payload, "rag_scope", None), + disable_parallel_tool_use = _disable_parallel, ) if payload.stream: @@ -4721,11 +5986,16 @@ async def anthropic_messages( _run_tool_gen, message_id, model_name, + llama_backend = llama_backend, + openai_messages = openai_messages, + openai_tools = openai_tools, + disable_parallel_tool_use = _disable_parallel, ) return await _anthropic_tool_non_streaming( _run_tool_gen, message_id, model_name, + disable_parallel_tool_use = _disable_parallel, ) # ── No-tool path ────────────────────────────────────────── @@ -4750,6 +6020,8 @@ async def anthropic_messages( _run_plain_gen, message_id, model_name, + llama_backend = llama_backend, + openai_messages = openai_messages, ) return await _anthropic_plain_non_streaming( _run_plain_gen, @@ -4764,15 +6036,38 @@ async def _anthropic_tool_stream( run_gen, message_id, model_name, + llama_backend = None, + openai_messages = None, + openai_tools = None, + disable_parallel_tool_use = False, ): """Streaming response for the tool-calling path.""" _sentinel = object() + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens + # makes blocking HTTP calls to llama-server, so run it off the event loop. + # Pass the tools so tool-schema tokens are counted (the generator renders + # them too), matching the non-stream / count_tokens / passthrough paths. + input_tokens = 0 + if llama_backend is not None and openai_messages is not None: + input_tokens = await asyncio.to_thread( + llama_backend.count_chat_tokens, openai_messages, None, openai_tools + ) + async def _stream(): emitter = AnthropicStreamEmitter() - for line in emitter.start(message_id, model_name): + for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line + captured_finish_reason = None + # Whether the response currently ends on a pending tool_use block (the + # client must act → stop_reason "tool_use") as opposed to final text. + # The server may run a tool and then keep generating, which flips this + # back to False — that is an end_turn (or max_tokens) response. + ends_on_tool_use = False + tool_blocks_emitted = 0 + drop_until_tool_end = False + gen = run_gen() try: while True: @@ -4782,16 +6077,53 @@ async def _anthropic_tool_stream( event = await asyncio.to_thread(next, gen, _sentinel) if event is _sentinel: break - # Strip leaked tool-call XML from content events - if event.get("type") == "content": + etype = event.get("type") + if drop_until_tool_end: + # disable_parallel_tool_use: a later tool call is being + # dropped — skip every event until (and including) its tool_end. + if etype == "tool_end": + drop_until_tool_end = False + continue + if etype == "metadata": + _fr = event.get("finish_reason") + if _fr is not None: + captured_finish_reason = _fr + # Strip leaked tool-call XML from content events first, so a + # content event that was purely tool XML doesn't count as text. + if etype == "content": event = dict(event) event["text"] = _TOOL_XML_RE.sub("", event["text"]) + # disable_parallel_tool_use: keep only the first tool_use block, + # dropping every later tool_start and its paired tool_end (robust + # to empty tool-call ids — tracked by state, not id matching). + if etype == "tool_start": + if disable_parallel_tool_use and tool_blocks_emitted >= 1: + drop_until_tool_end = True + continue + ends_on_tool_use = True + elif etype == "tool_end": + tool_blocks_emitted += 1 + # A tool_end means Studio executed the tool server-side, so + # the response no longer ends on a pending client action. + # Without this, a server tool that produces no trailing text + # would be mislabeled stop_reason "tool_use", telling the + # client to run a tool Studio already ran. + ends_on_tool_use = False + elif etype == "content" and event.get("text"): + ends_on_tool_use = False for line in emitter.feed(event): yield line except Exception as e: logger.error("anthropic_messages stream error: %s", e) + _error_event = _anthropic_stream_error_event(e) + if _error_event is not None: + yield _error_event + return - for line in emitter.finish("end_turn"): + stop_reason = openai_finish_to_anthropic_stop( + captured_finish_reason, had_tool_calls = ends_on_tool_use + ) + for line in emitter.finish(stop_reason = stop_reason, stop_sequence = None): yield line return StreamingResponse( @@ -4811,15 +6143,25 @@ async def _anthropic_plain_stream( run_gen, message_id, model_name, + llama_backend = None, + openai_messages = None, ): """Streaming response for the no-tool path.""" _sentinel = object() + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens + # makes blocking HTTP calls to llama-server, so run it off the event loop. + input_tokens = 0 + if llama_backend is not None and openai_messages is not None: + input_tokens = await asyncio.to_thread(llama_backend.count_chat_tokens, openai_messages) + async def _stream(): emitter = AnthropicStreamEmitter() - for line in emitter.start(message_id, model_name): + for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line + captured_finish_reason = None + gen = run_gen() try: while True: @@ -4831,6 +6173,9 @@ async def _anthropic_plain_stream( break if isinstance(cumulative, dict): if cumulative.get("type") == "metadata": + _fr = cumulative.get("finish_reason") + if _fr is not None: + captured_finish_reason = _fr for line in emitter.feed(cumulative): yield line continue @@ -4839,8 +6184,13 @@ async def _anthropic_plain_stream( yield line except Exception as e: logger.error("anthropic_messages stream error: %s", e) + _error_event = _anthropic_stream_error_event(e) + if _error_event is not None: + yield _error_event + return - for line in emitter.finish("end_turn"): + stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False) + for line in emitter.finish(stop_reason = stop_reason, stop_sequence = None): yield line return StreamingResponse( @@ -4854,25 +6204,62 @@ async def _anthropic_plain_stream( ) -async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): +def _anthropic_map_generation_error(e: Exception) -> HTTPException: + """Map an upstream 4xx / context-overflow generation error to a clean + Anthropic 400 invalid_request_error. Genuine 5xx errors stay 500.""" + if _classify_llama_generation_error(e) is not None: + return HTTPException( + status_code = 400, + detail = anthropic_error_body( + _friendly_error(e), + status = 400, + err_type = "invalid_request_error", + ), + ) + return HTTPException(status_code = 500, detail = _friendly_error(e)) + + +def _collect_anthropic_events(run_gen) -> list: + """Drain the generator into a list, mapping an upstream 4xx / context + overflow to a clean Anthropic 400 instead of leaking a 500.""" + try: + return list(run_gen()) + except HTTPException: + raise + except Exception as e: + raise _anthropic_map_generation_error(e) + + +async def _anthropic_tool_non_streaming( + run_gen, + message_id, + model_name, + disable_parallel_tool_use = False, +): """Non-streaming response for the tool-calling path. Builds ``content_blocks`` in generation order (text → tool_use → text → - tool_use → ...), mirroring the streaming emitter's behavior. Deltas - within a single synthesis turn are merged into the trailing text block; - tool_use blocks interrupt the text sequence and open a new text block on - the next content event. + tool_use → ...), mirroring the streaming emitter. Deltas within one + synthesis turn merge into the trailing text block; tool_use blocks interrupt + the text sequence and open a new text block on the next content event. ``prev_text`` is reset on ``tool_end`` because ``generate_chat_completion_with_tools`` yields cumulative content *per - turn* — the first content event of turn N+1 must diff against an empty - baseline, not against turn N's final length. + turn* -- the first content event of turn N+1 must diff against an empty + baseline, not turn N's final length. """ content_blocks: list = [] + tool_blocks_by_id: dict[str, AnthropicResponseToolUseBlock] = {} usage = {} prev_text = "" + captured_finish_reason = None + # Pending client tool_use; cleared by tool_end (server execution) or + # trailing text. See the stop_reason mapping below. + ends_on_tool_use = False - for event in run_gen(): + events = _collect_anthropic_events(run_gen) + + for event in events: etype = event.get("type", "") if etype == "content": # Strip leaked tool-call XML @@ -4880,30 +6267,66 @@ async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): new = clean[len(prev_text) :] prev_text = clean if new: - if content_blocks and isinstance( - content_blocks[-1], AnthropicResponseTextBlock - ): + ends_on_tool_use = False + if content_blocks and isinstance(content_blocks[-1], AnthropicResponseTextBlock): content_blocks[-1].text += new else: content_blocks.append(AnthropicResponseTextBlock(text = new)) elif etype == "tool_start": - content_blocks.append( - AnthropicResponseToolUseBlock( - id = event["tool_call_id"], + tool_call_id = event["tool_call_id"] + arguments = event.get("arguments", {}) + existing_tool_block = tool_blocks_by_id.get(tool_call_id) if tool_call_id else None + if existing_tool_block is not None: + if arguments or not existing_tool_block.input: + existing_tool_block.input = arguments + if event.get("tool_name") and not existing_tool_block.name: + existing_tool_block.name = event["tool_name"] + else: + tool_block = AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(tool_call_id), name = event["tool_name"], - input = event.get("arguments", {}), + input = arguments, ) - ) + if tool_call_id: + tool_blocks_by_id[tool_call_id] = tool_block + content_blocks.append(tool_block) + ends_on_tool_use = True elif etype == "tool_end": prev_text = "" + # Server-executed: no longer pending a client action (see above). + ends_on_tool_use = False elif etype == "metadata": usage = event.get("usage", {}) + _fr = event.get("finish_reason") + if _fr is not None: + captured_finish_reason = _fr + + # disable_parallel_tool_use: cap the response to at most one tool_use + # block. Keep the first tool_use and drop any later ones. + if disable_parallel_tool_use: + _seen_tool_use = False + _capped: list = [] + for block in content_blocks: + if isinstance(block, AnthropicResponseToolUseBlock): + if _seen_tool_use: + continue + _seen_tool_use = True + _capped.append(block) + content_blocks = _capped + + # stop_reason "tool_use" only when the response still ends on a pending + # tool_use (client must act). `ends_on_tool_use` is tracked through the + # event stream above: it is True only if the last tool_start had no + # following tool_end (server execution) or trailing text. + stop_reason = openai_finish_to_anthropic_stop( + captured_finish_reason, had_tool_calls = ends_on_tool_use + ) resp = AnthropicMessagesResponse( id = message_id, model = model_name, content = content_blocks, - stop_reason = "end_turn", + stop_reason = stop_reason, usage = AnthropicUsage( input_tokens = usage.get("prompt_tokens", 0), output_tokens = usage.get("completion_tokens", 0), @@ -4917,11 +6340,17 @@ async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): text_parts = [] usage = {} prev_text = "" + captured_finish_reason = None - for cumulative in run_gen(): + events = _collect_anthropic_events(run_gen) + + for cumulative in events: if isinstance(cumulative, dict): if cumulative.get("type") == "metadata": usage = cumulative.get("usage", {}) + _fr = cumulative.get("finish_reason") + if _fr is not None: + captured_finish_reason = _fr continue new = cumulative[len(prev_text) :] prev_text = cumulative @@ -4933,11 +6362,13 @@ async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): if full_text: content_blocks.append(AnthropicResponseTextBlock(text = full_text)) + stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False) + resp = AnthropicMessagesResponse( id = message_id, model = model_name, content = content_blocks, - stop_reason = "end_turn", + stop_reason = stop_reason, usage = AnthropicUsage( input_tokens = usage.get("prompt_tokens", 0), output_tokens = usage.get("completion_tokens", 0), @@ -4967,6 +6398,8 @@ def _build_passthrough_payload( response_format = None, chat_template_kwargs = None, backend_ctx = None, + seed = None, + stream_options = None, ): body = { "messages": openai_messages, @@ -4977,33 +6410,36 @@ def _build_passthrough_payload( "top_k": top_k, "stream": stream, } - if stream: - body["stream_options"] = {"include_usage": True} + if seed is not None: + body["seed"] = seed + if stream and stream_options is not None: + body["stream_options"] = stream_options body["max_tokens"] = ( - max_tokens - if max_tokens is not None - else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) + max_tokens if max_tokens is not None else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) ) body["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS - if stop: - body["stop"] = stop + # Normalize stop the same way the non-passthrough path does (the passthrough + # was previously the one path that forwarded an empty stop string verbatim). + _stop = _normalize_stop_sequences(stop) + if _stop: + body["stop"] = _stop if min_p is not None: body["min_p"] = min_p if repetition_penalty is not None: - # llama-server's field is "repeat_penalty", not "repetition_penalty" + # llama-server's field is "repeat_penalty", not "repetition_penalty". body["repeat_penalty"] = repetition_penalty if presence_penalty is not None: body["presence_penalty"] = presence_penalty if response_format is not None: - # llama-server applies a GBNF grammar derived from the JSON schema - # when response_format is present. Field is documented flat at the - # request root (tools/server/README.md), which is also what the - # OpenAI SDK produces by spreading extra_body into the body top. + # llama-server applies a GBNF grammar derived from the JSON schema when + # response_format is present. The field is documented flat at the + # request root (tools/server/README.md), which is also what the OpenAI + # SDK produces by spreading extra_body into the body top. body["response_format"] = response_format if chat_template_kwargs is not None: - # Propagate reasoning / template overrides (e.g. enable_thinking) - # so llama-server renders the Jinja template in the mode the caller - # asked for instead of whatever default the model was loaded with. + # Propagate reasoning / template overrides (e.g. enable_thinking) so + # llama-server renders the Jinja template in the caller's mode instead + # of the model's load-time default. body["chat_template_kwargs"] = chat_template_kwargs return body @@ -5027,9 +6463,10 @@ async def _anthropic_passthrough_stream( tool_choice = "auto", session_id = None, cancel_id = None, + disable_parallel_tool_use = False, ): """Streaming client-side pass-through: forward tools to llama-server and - translate its streaming response to Anthropic SSE without executing anything.""" + translate its stream to Anthropic SSE without executing anything.""" target_url = f"{llama_backend.base_url}/v1/chat/completions" body = _build_passthrough_payload( openai_messages, @@ -5045,6 +6482,15 @@ async def _anthropic_passthrough_stream( presence_penalty = presence_penalty, tool_choice = tool_choice, backend_ctx = llama_backend.context_length, + stream_options = {"include_usage": True}, + ) + + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens + # makes blocking HTTP calls to llama-server, so run it off the event loop. + # Pass the tools through so tool-schema tokens are counted (otherwise the + # streaming input_tokens undercounts vs the non-stream / count_tokens paths). + input_tokens = await asyncio.to_thread( + llama_backend.count_chat_tokens, openai_messages, None, openai_tools ) # cancel_id mirrors the OpenAI passthrough so a per-run cancel POST @@ -5054,36 +6500,32 @@ async def _anthropic_passthrough_stream( async def _stream(): emitter = AnthropicPassthroughEmitter() - for line in emitter.start(message_id, model_name): + for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line # Manage the httpx client, response, AND the aiter_lines() async - # generator MANUALLY — no `async with`, no anonymous iterator. + # generator MANUALLY -- no `async with`, no anonymous iterator. # # On Python 3.13 + httpcore 1.0.x, `async for raw_line in - # resp.aiter_lines():` creates an anonymous async generator. When - # the loop exits via `break` (or the generator is orphaned when a - # client disconnects mid-stream), Python's `async for` protocol - # does NOT auto-close the iterator the way a sync `for` loop - # would. The iterator remains reachable only from the current - # coroutine frame; once `_stream()` returns, the frame is GC'd - # and the iterator becomes unreachable. Python's asyncgen - # finalizer hook then runs its aclose() on a LATER GC pass in a - # DIFFERENT asyncio task, where httpcore's - # `HTTP11ConnectionByteStream.aclose()` enters - # `anyio.CancelScope.__exit__` with a mismatched task and prints - # `RuntimeError: Attempted to exit cancel scope in a different - # task` / `RuntimeError: async generator ignored GeneratorExit` - # as "Exception ignored in:" unraisable warnings. + # resp.aiter_lines():` creates an anonymous async generator. When the + # loop exits via `break` (or the generator is orphaned by a mid-stream + # client disconnect), `async for` does NOT auto-close the iterator like + # a sync `for` would. The iterator stays reachable only from the current + # coroutine frame; once `_stream()` returns, the frame is GC'd and the + # iterator becomes unreachable. The asyncgen finalizer then runs aclose() + # on a LATER GC pass in a DIFFERENT asyncio task, where httpcore's + # `HTTP11ConnectionByteStream.aclose()` enters `anyio.CancelScope.__exit__` + # with a mismatched task and prints `RuntimeError: Attempted to exit + # cancel scope in a different task` / `RuntimeError: async generator + # ignored GeneratorExit` as "Exception ignored in:" unraisable warnings. # - # The fix: save `resp.aiter_lines()` as `lines_iter`, and in the - # finally block explicitly `await lines_iter.aclose()` BEFORE - # `resp.aclose()` / `client.aclose()`. This closes the iterator - # inside our own task's event loop, so the internal httpcore - # byte-stream is cleaned up before Python's asyncgen finalizer - # has anything orphaned to finalize. Each aclose is wrapped in - # `try: ... except Exception: pass` so anyio cleanup noise from - # nested aclose paths can't bubble out. + # Fix: save `resp.aiter_lines()` as `lines_iter`, and in finally + # explicitly `await lines_iter.aclose()` BEFORE `resp.aclose()` / + # `client.aclose()`. This closes the iterator in our own task's event + # loop, cleaning up the httpcore byte-stream before the asyncgen + # finalizer has anything orphaned to finalize. Each aclose is wrapped in + # `try: ... except Exception: pass` so nested anyio cleanup noise can't + # bubble out. client = httpx.AsyncClient( timeout = 600, limits = httpx.Limits(max_keepalive_connections = 0), @@ -5095,13 +6537,32 @@ async def _anthropic_passthrough_stream( req = client.build_request("POST", target_url, json = body) resp = await client.send(req, stream = True) + # Upstream client error (e.g. over-context 400) arrives before any + # SSE. The 200 stream headers are already flushed, so surface it as + # an in-band Anthropic ``error`` event instead of silently finishing + # with an empty end_turn message. + if resp.status_code != 200: + _err_bytes = await resp.aread() + _err_text = _err_bytes.decode("utf-8", "replace")[:500] + logger.error( + "anthropic passthrough upstream error: status=%s body=%s", + resp.status_code, + _err_text, + ) + yield build_anthropic_sse_event( + "error", + anthropic_error_body( + f"llama-server error: {_err_text}", + status = resp.status_code, + ), + ) + return + # See _openai_passthrough_stream for rationale: aiter_lines() # blocks during llama-server prefill, so the in-loop cancel # check is unreachable until the first SSE chunk arrives. # The watcher closes `resp` on cancel, raising in aiter_lines. - cancel_watcher = asyncio.create_task( - _await_cancel_then_close(cancel_event, resp) - ) + cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) lines_iter = resp.aiter_lines() async for raw_line in lines_iter: if cancel_event.is_set(): @@ -5118,6 +6579,8 @@ async def _anthropic_passthrough_stream( chunk = json.loads(data_str) except json.JSONDecodeError: continue + if disable_parallel_tool_use: + _drop_parallel_tool_call_deltas(chunk) for line in emitter.feed_chunk(chunk): yield line except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): @@ -5177,6 +6640,7 @@ async def _anthropic_passthrough_non_streaming( repetition_penalty = None, presence_penalty = None, tool_choice = "auto", + disable_parallel_tool_use = False, ): """Non-streaming client-side pass-through.""" target_url = f"{llama_backend.base_url}/v1/chat/completions" @@ -5218,6 +6682,9 @@ async def _anthropic_passthrough_non_streaming( content_blocks.append(AnthropicResponseTextBlock(text = text)) tool_calls = message.get("tool_calls") or [] + # disable_parallel_tool_use: keep only the first tool_use block. + if disable_parallel_tool_use and len(tool_calls) > 1: + tool_calls = tool_calls[:1] for tc in tool_calls: fn = tc.get("function") or {} try: @@ -5226,18 +6693,13 @@ async def _anthropic_passthrough_non_streaming( args = {} content_blocks.append( AnthropicResponseToolUseBlock( - id = tc.get("id", ""), + id = anthropic_tool_use_id(tc.get("id")), name = fn.get("name", ""), input = args, ) ) - if tool_calls: - stop_reason = "tool_use" - elif finish_reason == "length": - stop_reason = "max_tokens" - else: - stop_reason = "end_turn" + stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls)) usage = data.get("usage") or {} resp_obj = AnthropicMessagesResponse( @@ -5271,24 +6733,122 @@ def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]: return out +_LOCAL_SERVER_BUILTIN_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "code_execution", "image_generation"} +) + + +def _strip_provider_synthetic_tool_history(messages: list[dict]) -> list[dict]: + """Drop synthetic provider-side tool_calls + matching role=tool replies on + the local-backend (llama-server / GGUF) dispatch path. + + A Gemini chat that ran code_execution / image_generation persists the + server-side tool card into history as an assistant tool_calls entry tagged + with ``args._server_tool`` (or a Gemini ``args.google.native_part`` payload) + plus a follow-up role=tool reply. When the user switches the SAME thread to + a local GGUF model, those synthetic tool_calls aren't real user functions, + llama-server has no matching declaration, and Gemini-only ``extra_content`` + / ``native_part`` payloads are meaningless. Forward only ordinary user + function calls; strip the matched role=tool replies too so the backend never + sees an orphan tool_call_id. + """ + dropped_ids: set[str] = set() + sanitized_assistant: list[dict] = [] + for m in messages: + if m.get("role") != "assistant": + sanitized_assistant.append(m) + continue + tool_calls = m.get("tool_calls") + if not isinstance(tool_calls, list) or not tool_calls: + # Plain text Gemini reply: still strip message-level + # `extra_content` (carries `google.thought_signature` replay + # metadata) so a text-only Gemini turn switched to a local GGUF + # backend doesn't leak Gemini-only fields to llama-server. + # ChatMessage didn't used to have `extra_content` (implicitly + # dropped); round-22 added it, which made this leak possible. + if "extra_content" in m: + m = {k: v for k, v in m.items() if k != "extra_content"} + sanitized_assistant.append(m) + continue + cleaned: list[dict] = [] + for tc in tool_calls: + if not isinstance(tc, dict): + cleaned.append(tc) + continue + fn = tc.get("function") + name = "" + if isinstance(fn, dict): + name = (fn.get("name") or "").lower() + if name in _LOCAL_SERVER_BUILTIN_TOOL_NAMES: + raw_args = fn.get("arguments") if isinstance(fn, dict) else None + args_obj: Any = None + if isinstance(raw_args, str): + try: + args_obj = json.loads(raw_args) if raw_args else None + except Exception: + args_obj = None + elif isinstance(raw_args, dict): + args_obj = raw_args + is_synthetic = False + if isinstance(args_obj, dict): + if args_obj.get("_server_tool") is True: + is_synthetic = True + google = args_obj.get("google") + if isinstance(google, dict) and isinstance(google.get("native_part"), dict): + is_synthetic = True + if is_synthetic: + tc_id = tc.get("id") + if isinstance(tc_id, str) and tc_id: + dropped_ids.add(tc_id) + continue + # Strip Gemini-only `extra_content` on real user tool_calls too -- + # llama-server has no use for it and may pass it to the model + # unchanged. + if "extra_content" in tc: + tc = {k: v for k, v in tc.items() if k != "extra_content"} + cleaned.append(tc) + # Drop message-level `extra_content` (Gemini thoughtSignature replay + # metadata) on local dispatch. + m_clean = {k: v for k, v in m.items() if k != "extra_content"} + if cleaned: + m_clean["tool_calls"] = cleaned + else: + m_clean.pop("tool_calls", None) + if not m_clean.get("content") and not m_clean.get("tool_calls"): + continue # assistant turn now empty, drop + sanitized_assistant.append(m_clean) + + if not dropped_ids: + return sanitized_assistant + out: list[dict] = [] + for m in sanitized_assistant: + if ( + m.get("role") == "tool" + and isinstance(m.get("tool_call_id"), str) + and m["tool_call_id"] in dropped_ids + ): + continue + out.append(m) + return out + + def _openai_messages_for_passthrough(payload) -> list[dict]: """Build OpenAI-format message dicts for the /v1/chat/completions passthrough path. - Messages from ``payload.messages`` are dumped through Pydantic (dropping - unset optional fields) so they are already in standard OpenAI format - — including ``role="tool"`` tool-result messages and assistant messages - that carry structured ``tool_calls``. Content-parts images already in - the message list are left untouched. + ``payload.messages`` are dumped through Pydantic (dropping unset optional + fields), so they're already standard OpenAI format -- including + ``role="tool"`` tool-result messages and assistant messages carrying + structured ``tool_calls``. Content-parts images already in the list are + left untouched. When a client uses Studio's legacy ``image_base64`` top-level field, the image is re-encoded to PNG (llama-server's stb_image has limited format - support) and spliced into the last user message as an OpenAI - ``image_url`` content part so vision + function-calling requests work - transparently. + support) and spliced into the last user message as an OpenAI ``image_url`` + content part so vision + function-calling requests work transparently. """ - messages = _drop_empty_assistant_sentinels( - [m.model_dump(exclude_none = True) for m in payload.messages] + messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels([m.model_dump(exclude_none = True) for m in payload.messages]) ) if not payload.image_base64: @@ -5333,12 +6893,12 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]: """Build llama-server messages for the standard GGUF chat path. - llama-server accepts OpenAI multimodal content parts directly. Preserve - all per-turn ``image_url`` parts so multi-image chat history keeps each - image attached to its original turn. + llama-server accepts OpenAI multimodal content parts directly. Preserve all + per-turn ``image_url`` parts so multi-image chat history keeps each image + attached to its original turn. """ - messages = _drop_empty_assistant_sentinels( - [m.model_dump(exclude_none = True) for m in payload.messages] + messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels([m.model_dump(exclude_none = True) for m in payload.messages]) ) has_message_image = any( isinstance(msg.get("content"), list) @@ -5373,10 +6933,10 @@ def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict] def _extract_response_format(payload): """Return the ``response_format`` field on an incoming ChatCompletionRequest - (or None). The model is declared with ``extra="allow"`` so pydantic stashes - unknown top-level fields in ``model_extra``; OpenAI-SDK clients spread - ``extra_body`` into the request body top level, which is where guided- - decoding recipes park their JSON-schema response_format. + (or None). The model uses ``extra="allow"`` so pydantic stashes unknown + top-level fields in ``model_extra``; OpenAI-SDK clients spread ``extra_body`` + into the request body top level, where guided-decoding recipes park their + JSON-schema response_format. """ extra = getattr(payload, "model_extra", None) if not isinstance(extra, dict): @@ -5385,28 +6945,42 @@ def _extract_response_format(payload): return rf if isinstance(rf, dict) else None -def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: +def _build_openai_passthrough_body( + payload, + backend_ctx = None, + llama_backend = None, +) -> dict: """Assemble the llama-server request body from a ChatCompletionRequest. - Only explicitly-known OpenAI / llama-server fields are forwarded so that - Studio-specific extensions (``enable_tools``, ``enabled_tools``, - ``session_id``, ...) never leak to the backend. + Only known OpenAI / llama-server fields are forwarded, so Studio-specific + extensions (``enable_tools``, ``enabled_tools``, ``session_id``, ...) never + leak to the backend. """ messages = _openai_messages_for_passthrough(payload) + system_prompt, _, _ = _extract_content_parts(payload.messages) + messages = _set_or_prepend_system_message(messages, system_prompt) tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto" - # When the caller asked for a specific reasoning mode, forward it to - # llama-server via chat_template_kwargs so the Jinja template renders - # with (or without) the reasoning preamble. - tpl_kwargs = None - if payload.enable_thinking is not None: - tpl_kwargs = {"enable_thinking": bool(payload.enable_thinking)} + # Forward per-request reasoning fields (enable_thinking / reasoning_effort / + # preserve_thinking) via chat_template_kwargs so the Jinja template renders + # in the caller's mode, gated on the active template's capabilities exactly + # like the non-passthrough paths. + tpl_kwargs = ( + llama_backend._request_reasoning_kwargs( + payload.enable_thinking, + payload.reasoning_effort, + payload.preserve_thinking, + ) + if llama_backend is not None + else None + ) return _build_passthrough_payload( messages, payload.tools, payload.temperature, payload.top_p, payload.top_k, - payload.max_tokens, + # Honor max_completion_tokens on the tools/response_format passthrough too. + _effective_max_tokens(payload), payload.stream, stop = payload.stop, min_p = payload.min_p, @@ -5416,68 +6990,69 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: response_format = _extract_response_format(payload), chat_template_kwargs = tpl_kwargs, backend_ctx = backend_ctx, + seed = payload.seed, + stream_options = payload.stream_options, ) async def _openai_passthrough_stream( - request, - cancel_event, - llama_backend, - payload, - model_name, - completion_id, + request, cancel_event, llama_backend, payload, model_name, completion_id ): """Streaming client-side pass-through for /v1/chat/completions. Forwards the client's OpenAI function-calling request to llama-server and - relays the SSE stream back verbatim. This preserves llama-server's - native response ``id``, ``finish_reason`` (including ``"tool_calls"``), - ``delta.tool_calls``, and the trailing ``usage`` chunk so the client - observes a standard OpenAI response. + relays the SSE stream back verbatim, preserving llama-server's native + response ``id``, ``finish_reason`` (including ``"tool_calls"``), + ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so + the client sees a standard OpenAI response. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" body = _build_openai_passthrough_body( - payload, backend_ctx = llama_backend.context_length + payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() - # Outer guard: asyncio.CancelledError at `await client.send(...)` is - # a BaseException that bypasses `except httpx.RequestError`; without - # this the tracker leaks. The generator's finally only runs once - # iteration starts. + # Outer guard: asyncio.CancelledError at `await client.send(...)` is a + # BaseException that bypasses `except httpx.RequestError`; without this the + # tracker leaks. The generator's finally only runs once iteration starts. try: - # Dispatch BEFORE returning StreamingResponse so transport errors - # and non-200 upstream statuses surface as real HTTP errors -- - # OpenAI SDKs rely on status codes to raise APIError/BadRequestError. + # Dispatch BEFORE returning StreamingResponse so transport errors and + # non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs + # rely on status codes to raise APIError/BadRequestError. client = httpx.AsyncClient( timeout = 600, limits = httpx.Limits(max_keepalive_connections = 0), ) resp = None - try: - req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) - except httpx.RequestError as e: - # llama-server subprocess crashed / still starting / unreachable. - logger.error("openai passthrough stream: upstream unreachable: %s", e) - if resp is not None: + _truncate_budget = ( + _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 + ) + while True: + try: + req = client.build_request("POST", target_url, json = body) + resp = await client.send(req, stream = True) + except httpx.RequestError as e: + # llama-server subprocess crashed / starting / unreachable. + logger.error("openai passthrough stream: upstream unreachable: %s", e) + if resp is not None: + try: + await resp.aclose() + except Exception: + pass try: - await resp.aclose() + await client.aclose() except Exception: pass - try: - await client.aclose() - except Exception: - pass - raise HTTPException( - status_code = 502, - detail = _friendly_error(e), - ) + raise HTTPException( + status_code = 502, + detail = _friendly_error(e), + ) - if resp.status_code != 200: + if resp.status_code == 200: + break err_bytes = await resp.aread() err_text = err_bytes.decode("utf-8", errors = "replace") logger.error( @@ -5490,30 +7065,32 @@ async def _openai_passthrough_stream( await resp.aclose() except Exception: pass + # Opt-in overflow policy: shrink and retry instead of a fatal 400. + if ( + _truncate_budget > 0 + and _classify_llama_generation_error(Exception(err_text)) + and _apply_overflow_truncation(body, err_text) + ): + _truncate_budget -= 1 + continue try: await client.aclose() except Exception: pass - raise HTTPException( - status_code = upstream_status, - detail = f"llama-server error: {err_text[:500]}", - ) + raise _openai_passthrough_error(upstream_status, err_text) async def _stream(): # Same httpx lifecycle pattern as _anthropic_passthrough_stream: - # save resp.aiter_lines() so the finally block can aclose() it - # on our task. See that function for full rationale. + # save resp.aiter_lines() so the finally block can aclose() it on + # our task. See that function for full rationale. lines_iter = None # During llama-server prefill, `aiter_lines()` blocks until the - # first SSE chunk arrives. The in-loop `cancel_event` check - # cannot fire until then, which is the exact proxy/Colab - # scenario the cancel POST is meant to recover from. Run a - # tiny watcher that closes `resp` as soon as cancel fires, - # unblocking the iterator with a RemoteProtocolError caught - # in the except clause below. - cancel_watcher = asyncio.create_task( - _await_cancel_then_close(cancel_event, resp) - ) + # first SSE chunk arrives. The in-loop `cancel_event` check can't + # fire until then -- the exact proxy/Colab scenario the cancel POST + # recovers from. Run a tiny watcher that closes `resp` as soon as + # cancel fires, unblocking the iterator with a RemoteProtocolError + # caught in the except clause below. + cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) try: lines_iter = resp.aiter_lines() async for raw_line in lines_iter: @@ -5526,25 +7103,26 @@ async def _openai_passthrough_stream( continue if not raw_line.startswith("data: "): continue + # Honor parallel_tool_calls=false (best-effort): drop tool_call + # deltas with index>=1 so only the first call streams. Only + # lines carrying tool_calls are reparsed; everything else is + # relayed byte-for-byte. + if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: + raw_line = _cap_parallel_tool_calls_sse_line(raw_line) # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. yield raw_line + "\n\n" if raw_line[6:].strip() == "[DONE]": break except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): - # Watcher closed resp on cancel. Emit nothing extra; the - # client either initiated the cancel or already disconnected. + # Watcher closed resp on cancel. Emit nothing extra; the client + # initiated the cancel or already disconnected. if not cancel_event.is_set(): raise except Exception as e: - # 200 headers are already flushed; errors must be in the SSE body. + # 200 headers already flushed; errors must go in the SSE body. logger.error("openai passthrough stream error: %s", e) - err = { - "error": { - "message": _friendly_error(e), - "type": "server_error", - }, - } + err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: cancel_watcher.cancel() @@ -5581,79 +7159,101 @@ async def _openai_passthrough_stream( raise -async def _openai_passthrough_non_streaming( - llama_backend, - payload, - model_name, -): +async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): """Non-streaming client-side pass-through for /v1/chat/completions. - Returns llama-server's JSON response verbatim (via JSONResponse) so the - client sees the native response ``id``, ``finish_reason`` (including - ``"tool_calls"``), structured ``tool_calls``, and accurate ``usage`` - token counts. + Returns llama-server's JSON response verbatim so the client sees the native + response ``id``, ``finish_reason`` (including ``"tool_calls"``), structured + ``tool_calls``, and accurate ``usage`` token counts. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" body = _build_openai_passthrough_body( - payload, backend_ctx = llama_backend.context_length + payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) - try: - async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) - except httpx.RequestError as e: - # llama-server subprocess crashed / still starting / unreachable. - # Surface the same friendly message the sync chat path emits so - # operators don't see a bare 500 with no diagnostic. - logger.error("openai passthrough non-streaming: upstream unreachable: %s", e) - raise HTTPException( - status_code = 502, - detail = _friendly_error(e), - ) - - if resp.status_code != 200: - raise HTTPException( - status_code = resp.status_code, - detail = f"llama-server error: {resp.text[:500]}", - ) - - # Guided-decoding fence wrap. llama-server returns raw JSON that matches - # the schema (no surrounding markdown) because the GBNF grammar only - # emits the JSON object itself. data_designer's llm-structured parser - # looks for a ```json ... ``` markdown fence and discards unfenced - # output, which collapses a 100%-valid guided-decoding run to 0/N. - # Wrap each choice's content in the expected fence when the caller - # asked for guided decoding, leaving already-fenced content alone. - if _extract_response_format(payload) is not None: + _truncate_budget = ( + _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 + ) + while True: try: - data = resp.json() - changed = False - for choice in data.get("choices", []): - if not isinstance(choice, dict): - continue - msg = choice.get("message") - if not isinstance(msg, dict): - continue - content = msg.get("content") - if not isinstance(content, str): - continue - stripped = content.strip() - if not stripped or stripped.startswith("```"): - continue - msg["content"] = f"```json\n{stripped}\n```" - changed = True - if changed: - return JSONResponse(content = data) - except Exception as exc: - # Wrap is best-effort; fall through to the verbatim body if - # the response is not JSON-shaped or the structure is unusual. - logger.warning( - "response_format fence wrap skipped: %s", - exc, + async with httpx.AsyncClient() as client: + resp = await client.post(target_url, json = body, timeout = 600) + except httpx.RequestError as e: + # llama-server subprocess crashed / starting / unreachable. Surface the + # same friendly message the sync chat path emits so operators don't see + # a bare 500 with no diagnostic. + logger.error("openai passthrough non-streaming: upstream unreachable: %s", e) + raise HTTPException( + status_code = 502, + detail = _friendly_error(e), ) - # Pass the upstream body through as raw bytes — skips a redundant - # parse+re-serialize round-trip and keeps the response truly - # verbatim (matches the docstring). Status is guaranteed 200 by - # the check above. - return Response(content = resp.content, media_type = "application/json") + if resp.status_code == 200: + break + # Opt-in overflow policy: shrink and retry instead of a fatal 400. + if ( + _truncate_budget > 0 + and _classify_llama_generation_error(Exception(resp.text)) + and _apply_overflow_truncation(body, resp.text) + ): + _truncate_budget -= 1 + continue + raise _openai_passthrough_error(resp.status_code, resp.text) + + # The guided-decoding fence wraps each choice's JSON content in a + # ```json ... ``` markdown fence that data_designer's structured parser + # requires but which CORRUPTS output for standard OpenAI clients doing + # ``json.loads(content)``. It is therefore opt-in: only the internal + # data-recipe path sets ``_unsloth_guided_fence``; public response_format + # clients get the raw upstream JSON verbatim. + _guided_fence = bool((payload.model_extra or {}).get("_unsloth_guided_fence")) + _do_fence = _guided_fence and _extract_response_format(payload) is not None + _cap_parallel = payload.parallel_tool_calls is False + + try: + data = resp.json() + except Exception as exc: + # Non-JSON / unparseable upstream body: relay verbatim as before. + logger.warning( + "openai passthrough non-streaming: response not JSON, relaying raw: %s", + exc, + ) + return Response(content = resp.content, media_type = "application/json") + + changed = False + for choice in data.get("choices", []): + if not isinstance(choice, dict): + continue + msg = choice.get("message") + if not isinstance(msg, dict): + continue + + # OpenAI requires content=null on a pure tool-call turn; llama-server + # emits content="". + if msg.get("tool_calls") and msg.get("content") == "": + msg["content"] = None + changed = True + + # Honor parallel_tool_calls=false (best-effort) by capping to one call. + if _cap_parallel: + _tcs = msg.get("tool_calls") + if isinstance(_tcs, list) and len(_tcs) > 1: + msg["tool_calls"] = _tcs[:1] + changed = True + + # Guided-decoding fence wrap (opt-in via _unsloth_guided_fence). + if _do_fence: + content = msg.get("content") + if not isinstance(content, str): + continue + stripped = content.strip() + if not stripped or stripped.startswith("```"): + continue + msg["content"] = f"```json\n{stripped}\n```" + changed = True + + # Nothing mutated: relay the upstream bytes verbatim, skipping a redundant + # parse + re-serialize round-trip. + if not changed: + return Response(content = resp.content, media_type = "application/json") + return JSONResponse(content = data) diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py new file mode 100644 index 0000000000..3aae6f4209 --- /dev/null +++ b/studio/backend/routes/llama.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""llama.cpp prebuilt update endpoints. + +GET /api/llama/update-status -> is a newer prebuilt available + job state +POST /api/llama/update -> download + atomically swap to the latest + +Detection reuses utils.llama_cpp_freshness; the swap reuses +install_llama_prebuilt.py via utils.llama_cpp_update. Both fail open so the UI +never blocks on a missing marker / offline GitHub. +""" + +from __future__ import annotations + +import asyncio +from typing import Optional + +from fastapi import APIRouter, Depends, Query +from pydantic import BaseModel, Field + +from auth.authentication import get_current_subject +from utils.llama_cpp_update import get_update_status, start_update + +router = APIRouter() + + +class LlamaUpdateJob(BaseModel): + state: str = Field("idle", description = "idle | running | success | error") + message: str = "" + from_tag: Optional[str] = None + to_tag: Optional[str] = None + error: Optional[str] = None + progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.") + started_at: Optional[str] = None + finished_at: Optional[str] = None + + +class LlamaUpdateStatusResponse(BaseModel): + supported: bool = Field( + False, + description = "True when the install came from an Unsloth prebuilt (has a marker).", + ) + update_available: bool = Field( + False, description = "True when the latest release is genuinely newer than the install." + ) + stale: bool = Field( + False, description = "Update available AND install older than the staleness threshold." + ) + installed_tag: Optional[str] = None + latest_tag: Optional[str] = None + published_repo: Optional[str] = None + installed_at_utc: Optional[str] = None + age_days: Optional[int] = None + source_build: bool = Field( + False, description = "True when there is no marker (source build) but a prebuilt is offered." + ) + job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) + + +class LlamaUpdateActionResponse(BaseModel): + started: bool + reason: Optional[str] = None + message: Optional[str] = None + job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) + + +@router.get("/update-status", response_model = LlamaUpdateStatusResponse) +async def llama_update_status( + force_refresh: bool = Query( + False, description = "Bypass the 24h release cache for an explicit check." + ), + current_subject: str = Depends(get_current_subject), +) -> LlamaUpdateStatusResponse: + # Off the event loop: detection may probe the host and read GitHub. + status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh) + return LlamaUpdateStatusResponse(**status) + + +@router.post("/update", response_model = LlamaUpdateActionResponse) +async def llama_update( + current_subject: str = Depends(get_current_subject), +) -> LlamaUpdateActionResponse: + action = await asyncio.to_thread(start_update) + return LlamaUpdateActionResponse(**action) diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py new file mode 100644 index 0000000000..3001c6b7c9 --- /dev/null +++ b/studio/backend/routes/mcp_servers.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import json +import uuid +from urllib.parse import urlparse + +import structlog +from fastapi import APIRouter, Depends, HTTPException + +from auth.authentication import get_current_subject +from core.inference.mcp_client import ( + clear_oauth_tokens_async, + is_stdio, + list_tools_async, + parse_server_headers, + parse_stdio_command, + probe_timeout, + stdio_mcp_enabled, +) +from core.inference.mcp_config_import import parse_mcp_config +from models.mcp_servers import ( + McpServerCreate, + McpServerImportRequest, + McpServerImportResult, + McpServerProbeResult, + McpServerResponse, + McpServerTestRequest, + McpServerUpdate, +) +from storage import mcp_servers_db +from utils.utils import safe_curated_detail, log_and_http_error + +logger = structlog.get_logger(__name__) + +router = APIRouter() + + +def _looks_like_command(value: str) -> bool: + """Whitespace is a one-way signal: a URL can't hold an unencoded space, so + a value with whitespace is definitely a command. No whitespace proves + nothing (a lone token may be a single-arg command or a scheme-less URL).""" + return any(ch.isspace() for ch in value) + + +def _validate_url(url: str) -> str: + trimmed = (url or "").strip() + if not trimmed: + raise HTTPException(status_code = 400, detail = "url must not be empty") + # When stdio is enabled, a non-HTTP value is a local command (reuses this + # field so stdio servers ride existing CRUD/storage). + if stdio_mcp_enabled() and is_stdio(trimmed): + try: + parts = parse_stdio_command(trimmed) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + "Invalid command. Check quoting and try again.", + event = "mcp_servers.invalid_command", + log = logger, + ) + if not parts or not parts[0].strip(): + raise HTTPException(status_code = 400, detail = "command must not be empty") + if "://" in parts[0]: + # A URL-scheme first token is a mistyped URL, not a command. Reject + # cleanly instead of exec-ing it (mirrors the frontend check). + raise HTTPException( + status_code = 400, + detail = "Enter an http(s):// URL, or a local command whose " + "first token is an executable (not a URL).", + ) + return trimmed + parsed = urlparse(trimmed) + if parsed.scheme not in ("http", "https"): + detail = ( + "MCP server address must start with http:// or https:// " + "(for example https://example.com/mcp)." + ) + # Host-scoped wording: self-hosted hosts can opt in via the env var. + if _looks_like_command(trimmed): + detail += " Running a local command is not enabled on this server." + raise HTTPException(status_code = 400, detail = detail) + if not parsed.netloc: + raise HTTPException(status_code = 400, detail = "url is missing a host") + return trimmed + + +def _normalize_headers(headers: dict[str, str] | None) -> dict[str, str] | None: + """Trim header names, drop empties, coerce values to str; None if empty.""" + if not headers: + return None + out: dict[str, str] = {} + for raw_key, value in headers.items(): + key = str(raw_key).strip() + if key: + out[key] = str(value) + return out or None + + +def _row_to_response(row: dict) -> McpServerResponse: + return McpServerResponse( + id = row["id"], + display_name = row["display_name"], + url = row["url"], + headers = parse_server_headers(row) or {}, + is_enabled = bool(row["is_enabled"]), + use_oauth = bool(row.get("use_oauth")), + created_at = row["created_at"], + updated_at = row["updated_at"], + ) + + +@router.get("/", response_model = list[McpServerResponse]) +async def list_mcp_servers(current_subject: str = Depends(get_current_subject)): + return [_row_to_response(row) for row in mcp_servers_db.list_servers()] + + +@router.post("/", response_model = McpServerResponse, status_code = 201) +async def create_mcp_server( + payload: McpServerCreate, current_subject: str = Depends(get_current_subject) +): + display_name = (payload.display_name or "").strip() + if not display_name: + raise HTTPException(status_code = 400, detail = "display_name must not be empty") + url = _validate_url(payload.url) + headers = _normalize_headers(payload.headers) + # OAuth is HTTP-only; force it off for stdio commands so a stale flag can't + # push the probe onto the 305s OAuth timeout. Backend enforces this. + use_oauth = payload.use_oauth and not is_stdio(url) + + server_id = uuid.uuid4().hex[:16] + mcp_servers_db.create_server( + id = server_id, + display_name = display_name, + url = url, + headers_json = json.dumps(headers) if headers else None, + is_enabled = payload.is_enabled, + use_oauth = use_oauth, + ) + return _row_to_response(mcp_servers_db.get_server(server_id)) + + +def _changes_from_payload(payload: McpServerUpdate) -> dict: + sent = payload.model_fields_set + changes: dict = {} + + if "display_name" in sent: + name = (payload.display_name or "").strip() + if not name: + raise HTTPException(status_code = 400, detail = "display_name must not be empty") + changes["display_name"] = name + if "url" in sent: + changes["url"] = _validate_url(payload.url or "") + if "headers" in sent: + headers = _normalize_headers(payload.headers) + changes["headers_json"] = json.dumps(headers) if headers else None + if "is_enabled" in sent: + if payload.is_enabled is None: + raise HTTPException(status_code = 400, detail = "is_enabled must be true or false") + changes["is_enabled"] = payload.is_enabled + if "use_oauth" in sent: + if payload.use_oauth is None: + raise HTTPException(status_code = 400, detail = "use_oauth must be true or false") + changes["use_oauth"] = payload.use_oauth + # stdio is OAuth-less: drop a stale OAuth flag when switching to a command. + if "url" in changes and is_stdio(changes["url"]): + changes["use_oauth"] = False + return changes + + +@router.put("/{server_id}", response_model = McpServerResponse) +async def update_mcp_server( + server_id: str, + payload: McpServerUpdate, + current_subject: str = Depends(get_current_subject), +): + old = mcp_servers_db.get_server(server_id) + if not old: + raise HTTPException(status_code = 404, detail = "MCP server not found") + changes = _changes_from_payload(payload) + if not changes: + raise HTTPException(status_code = 400, detail = "No fields to update") + # headers == HTTP headers (remote) or env vars (stdio). On a transport-type + # switch with no new headers, drop the old ones so env secrets aren't + # re-sent as HTTP headers (or vice versa). + if ( + "url" in changes + and is_stdio(changes["url"]) != is_stdio(old["url"]) + and "headers_json" not in changes + ): + changes["headers_json"] = None + # Clear persisted OAuth tokens when the URL changes or OAuth is disabled; + # fastmcp keys tokens by URL and would otherwise let a re-pointed server + # silently inherit the old account's credentials. + if bool(old.get("use_oauth")) and ( + ("url" in changes and changes["url"] != old["url"]) or changes.get("use_oauth") is False + ): + await clear_oauth_tokens_async(old["url"]) + mcp_servers_db.update_server(server_id, changes) + return _row_to_response(mcp_servers_db.get_server(server_id)) + + +@router.delete("/{server_id}", status_code = 204) +async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_current_subject)): + old = mcp_servers_db.get_server(server_id) + if not old: + raise HTTPException(status_code = 404, detail = "MCP server not found") + if old.get("use_oauth"): + await clear_oauth_tokens_async(old["url"]) + mcp_servers_db.delete_server(server_id) + + +@router.post("/{server_id}/refresh", response_model = McpServerProbeResult) +async def refresh_mcp_server_tools( + server_id: str, current_subject: str = Depends(get_current_subject) +): + server = mcp_servers_db.get_server(server_id) + if not server: + raise HTTPException(status_code = 404, detail = "MCP server not found") + # Refresh uses the stored address, so re-check the stdio gate here too: a + # stdio row from a desktop DB must not spawn on a hosted/network host. + if is_stdio(server["url"]) and not stdio_mcp_enabled(): + raise HTTPException(status_code = 400, detail = "stdio MCP servers are disabled on this host") + + use_oauth = bool(server.get("use_oauth")) + try: + tools = await list_tools_async( + url = server["url"], + headers = parse_server_headers(server), + timeout = probe_timeout(server["url"], use_oauth), + use_oauth = use_oauth, + ) + except Exception as exc: # noqa: BLE001 — surface transport+timeout errors to UI + logger.error( + "mcp_servers.refresh_failed", + server_id = server_id, + error = str(exc), + exc_info = True, + ) + return McpServerProbeResult(ok = False, error = safe_curated_detail(exc)) + + return McpServerProbeResult(ok = True, tool_count = len(tools)) + + +@router.post("/import", response_model = McpServerImportResult) +async def import_mcp_servers( + payload: McpServerImportRequest, current_subject: str = Depends(get_current_subject) +): + """Bulk-register servers from a standard mcpServers JSON config (issue + #5936). Each entry rides the existing create path: _validate_url applies + the same stdio gate (a stdio entry becomes a per-entry error when stdio is + off; http still imports), and entries whose url already exists are skipped + so re-importing the same file is idempotent. One bad entry never 400s the + whole batch -- failures are reported per entry.""" + entries, errors = parse_mcp_config(payload.config) + created: list[McpServerResponse] = [] + skipped: list[str] = [] + seen_urls = {row["url"] for row in mcp_servers_db.list_servers()} + + for entry in entries: + try: + url = _validate_url(entry.url) + except HTTPException as exc: + errors.append(f"{entry.display_name}: {exc.detail}") + continue + if url in seen_urls: + skipped.append(entry.display_name) + continue + headers = _normalize_headers(entry.headers) + server_id = uuid.uuid4().hex[:16] + mcp_servers_db.create_server( + id = server_id, + display_name = entry.display_name, + url = url, + headers_json = json.dumps(headers) if headers else None, + is_enabled = entry.is_enabled, + use_oauth = entry.use_oauth and not is_stdio(url), + ) + seen_urls.add(url) + created.append(_row_to_response(mcp_servers_db.get_server(server_id))) + + return McpServerImportResult(created = created, skipped = skipped, errors = errors) + + +@router.post("/test", response_model = McpServerProbeResult) +async def test_mcp_server( + payload: McpServerTestRequest, current_subject: str = Depends(get_current_subject) +): + # URL/header validation must surface as 400 like create/update so the + # frontend's create-form pre-flight gets the same error semantics as the + # save call. Only catch transport/timeout errors below. + url = _validate_url(payload.url) + headers = _normalize_headers(payload.headers) + try: + tools = await list_tools_async( + url = url, + headers = headers, + timeout = probe_timeout(url, payload.use_oauth), + use_oauth = payload.use_oauth, + ) + except Exception as exc: # noqa: BLE001 + logger.error( + "mcp_servers.test_failed", + error = str(exc), + exc_info = True, + ) + return McpServerProbeResult(ok = False, error = safe_curated_detail(exc)) + + return McpServerProbeResult(ok = True, tool_count = len(tools)) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 9ea113e488..7151af33f3 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Model Management API routes -""" +"""Model management API routes.""" import hashlib import json @@ -16,6 +14,7 @@ from fastapi import APIRouter, Body, Depends, HTTPException, Query from typing import List, Optional import structlog from loggers import get_logger +from utils.utils import log_and_http_error import re as _re @@ -27,14 +26,12 @@ def _is_valid_repo_id(repo_id: str) -> bool: def _safe_is_dir(path) -> bool: - """``Path.is_dir()`` that returns ``False`` instead of raising. + """``Path.is_dir()`` returning ``False`` instead of raising. - On Python >= 3.12 ``is_dir()``'s ``os.stat`` only suppresses - "not found"-class errors and now propagates ``PermissionError`` - (EACCES); on Python <= 3.11 it returned ``False``. The folder-scan - endpoints probe well-known system locations (e.g. a root-owned, - mode-700 ``/usr/share/ollama/.ollama/models``) and must treat an - un-stat-able path as "not a directory", never 500. + Python >= 3.12 propagates ``PermissionError`` from ``is_dir()``; + folder-scan endpoints probe system locations (e.g. root-owned + ``/usr/share/ollama``) and must treat un-stat-able paths as "not a + directory", never 500. """ try: return Path(path).is_dir() @@ -42,14 +39,24 @@ def _safe_is_dir(path) -> bool: return False -# Add backend directory to path +def _is_hidden_model(*values: str | None) -> bool: + """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or + EMBED_GGUF_REPO basename), so pickers hide it (GGUF and non-GGUF).""" + from core.rag import config as rag_config + + needles = ( + rag_config.EMBEDDING_MODEL.split("/")[-1].lower(), + rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(), + ) + return any(v and any(n in v.lower() for n in needles) for v in values) + + backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) from auth.authentication import get_current_subject -# Import backend functions try: from utils.models import ( scan_trained_models, @@ -78,7 +85,7 @@ try: resolve_export_dir, ) except ImportError: - # Fallback: try to import from parent directory + # Fallback: import from parent directory. parent_backend = backend_path.parent / "backend" if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) @@ -140,7 +147,9 @@ logger = get_logger(__name__) def derive_model_type( - is_vision: bool, audio_type: Optional[str], is_embedding: bool = False + is_vision: bool, + audio_type: Optional[str], + is_embedding: bool = False, ) -> ModelType: """Collapse individual capability flags into a single model modality string.""" if is_embedding: @@ -156,7 +165,6 @@ def _resolve_hf_cache_dir() -> Path: """Resolve local HF cache root used by hub downloads.""" try: from huggingface_hub.constants import HF_HUB_CACHE - return Path(HF_HUB_CACHE) except Exception: return Path.home() / ".cache" / "huggingface" / "hub" @@ -165,15 +173,10 @@ def _resolve_hf_cache_dir() -> Path: def _is_model_directory(d: Path) -> bool: """Return ``True`` when *d* looks like a model directory. - A model directory must have **both** a config file (``config.json`` or - ``adapter_config.json``) **and** actual model weight files. Both - conditions are required: a bare directory with only loose ``.gguf`` - files (no config) might be a mixed collection, and a ``config.json`` - alone (no weights) is not a model directory. - - Excludes ``mmproj`` GGUF files (vision projectors) and non-weight - ``.bin`` files (``tokenizer.bin``, ``vocab.bin``, etc.) from the - weight check to avoid false positives. + Requires both a config (``config.json``/``adapter_config.json``) and + weight files. Excludes ``mmproj`` GGUFs (vision projectors) and + non-weight ``.bin`` files (``tokenizer.bin`` etc.) to avoid false + positives. """ def _is_weight_file(f: Path) -> bool: @@ -193,9 +196,7 @@ def _is_model_directory(d: Path) -> bool: return False try: - has_config = (d / "config.json").exists() or ( - d / "adapter_config.json" - ).exists() + has_config = (d / "config.json").exists() or (d / "adapter_config.json").exists() if not has_config: return False return any(_is_weight_file(f) for f in d.iterdir() if f.is_file()) @@ -203,11 +204,7 @@ def _is_model_directory(d: Path) -> bool: return False -def _scan_models_dir( - models_dir: Path, - *, - limit: int | None = None, -) -> List[LocalModelInfo]: +def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[LocalModelInfo]: if not models_dir.exists() or not models_dir.is_dir(): return [] @@ -243,8 +240,7 @@ def _scan_models_dir( or any(child.glob("*.gguf")) ) except OSError: - # Skip individual children that are unreadable (permissions, broken - # symlinks, etc.) rather than failing the entire scan. + # Skip unreadable children rather than failing the scan. continue if not has_model_files: continue @@ -261,7 +257,7 @@ def _scan_models_dir( updated_at = updated_at, ), ) - # Also scan for standalone .gguf files directly in the models directory + # Also scan standalone .gguf files in the models directory. if limit is None or len(found) < limit: for gguf_file in models_dir.glob("*.gguf"): if limit is not None and len(found) >= limit: @@ -319,16 +315,14 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]: def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: """Scan an LM Studio models directory for model files. - LM Studio uses a ``publisher/model-name`` folder structure containing - GGUF files, or standalone GGUF files at the top level. + LM Studio uses a ``publisher/model-name`` folder structure with GGUF + files, or standalone GGUF files at the top level. """ if not lm_dir.exists() or not lm_dir.is_dir(): return [] - # If the directory itself is a model directory (has config AND weight - # files), it is not an LM Studio publisher structure -- return it as a - # single model entry. We cannot skip it silently because this function - # is the only scanner called for default LM Studio roots. + # If lm_dir is itself a model directory (not a publisher structure), + # return it as a single entry rather than skipping it silently. if _is_model_directory(lm_dir): try: updated_at = lm_dir.stat().st_mtime @@ -364,9 +358,8 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: ) continue - # If the child directory itself looks like a model directory - # (has config AND weight files), surface it directly instead - # of descending into it as a publisher. + # Surface a model-directory child directly instead of + # descending into it as a publisher. if _is_model_directory(child): try: updated_at = child.stat().st_mtime @@ -383,7 +376,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: ) continue - # child is a publisher directory -- scan its sub-directories + # child is a publisher directory; scan its subdirectories. for model_dir in child.iterdir(): try: if model_dir.is_dir(): @@ -434,11 +427,9 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: """Return a writable directory for Ollama ``.gguf`` symlinks. - Prefers ``/.studio_links/`` so the links sit next to the - blobs they point at. Falls back to a per-ollama-dir namespace under - Studio's own cache when the models directory is read-only (common - for system installs under ``/usr/share/ollama`` or ``/var/lib/ollama``) - so we still surface Ollama models in those environments. + Prefers ``/.studio_links/`` so links sit next to their + blobs; falls back to a per-ollama-dir namespace under Studio's cache + when the models dir is read-only (common for system installs). """ from utils.paths.storage_roots import cache_root @@ -448,15 +439,13 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: return primary except OSError as e: logger.debug( - "Ollama dir %s not writable for .studio_links (%s); " - "falling back to Studio cache", + "Ollama dir %s not writable for .studio_links (%s); falling back to Studio cache", ollama_dir, e, ) - # Fallback: namespace by a hash of the ollama_dir so two different - # Ollama roots don't collide. This is a cache path, not a security - # boundary. + # Fallback: namespace by a hash of ollama_dir so two roots don't + # collide. Cache path, not a security boundary. try: digest = hashlib.sha256(str(ollama_dir.resolve()).encode()).hexdigest()[:12] except OSError: @@ -474,39 +463,22 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: return None -def _scan_ollama_dir( - ollama_dir: Path, limit: Optional[int] = None -) -> List[LocalModelInfo]: +def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[LocalModelInfo]: """Scan an Ollama models directory for downloaded models. - Ollama stores models in a content-addressable layout:: + Ollama uses a content-addressable layout + (``manifests////`` + ``blobs/sha256-...``); + we ``rglob`` all manifests so every layout depth is found. Each + manifest is JSON with a ``layers`` array: the + ``application/vnd.ollama.image.model`` layer holds the GGUF weights + and ``...image.projector`` is the vision adapter. - /manifests//// - /blobs/sha256-... - - The default host is ``registry.ollama.ai`` with namespace - ``library`` (official models), but users can pull from custom - namespaces (``mradermacher/llama3``) or entirely different hosts - (``hf.co/org/repo:tag``). We iterate all manifest files via - ``rglob`` so every layout depth is discovered. - - Each manifest is JSON with a ``layers`` array. The layer with - ``mediaType == "application/vnd.ollama.image.model"`` contains the - GGUF weights. Vision models also have a projector layer - (``application/vnd.ollama.image.projector``). We read the config - layer to extract family/size info. - - Since Ollama blobs lack a ``.gguf`` extension (which the GGUF - loading pipeline requires), we create ``.gguf``-named links - pointing at the blobs so the existing ``detect_gguf_model`` and - ``llama-server -m`` paths work unchanged. Each model gets its - own subdirectory under the links dir (keyed by a short hash of - the manifest path) so that ``detect_mmproj_file`` only sees the - projector for *that* model. Links are created as symlinks when - possible, falling back to hardlinks (Windows without Developer - Mode) as a last resort. The link dir lives under - ``/.studio_links/`` when writable, otherwise under - Studio's own cache directory. + Ollama blobs lack the ``.gguf`` extension the loading pipeline + requires, so we create ``.gguf``-named links to them (one subdir per + model, keyed by a short hash of the manifest path, so + ``detect_mmproj_file`` only sees that model's projector). Links are + symlinks when possible, else hardlinks; the link dir is + ``.studio_links/`` when writable, else Studio's cache. """ manifests_root = ollama_dir / "manifests" if not manifests_root.is_dir(): @@ -525,20 +497,16 @@ def _scan_ollama_dir( def _make_link(link_dir: Path, link_name: str, target: Path) -> Optional[str]: """Create a .gguf-named link to an Ollama blob. - Tries symlink first, then hardlink (works on Windows without - Developer Mode when target is on the same filesystem). Skips - the model if neither works -- a full file copy of a multi-GB - GGUF inside a synchronous API request would block the backend. - + Tries symlink, then hardlink; skips the model if neither works + (a multi-GB copy in a sync request would block the backend). Idempotent: skips recreation when a valid link already exists. """ link_dir.mkdir(parents = True, exist_ok = True) link_path = link_dir / link_name resolved = target.resolve() - # Skip if the link already points at the exact same blob. - # Only use samefile -- size-based checks can reuse stale links - # after `ollama pull` updates a tag to a same-sized blob. + # Skip if the link already points at the same blob. Use samefile + # only; size checks can reuse stale links after `ollama pull`. try: if link_path.exists() and os.path.samefile(str(link_path), str(resolved)): return str(link_path) @@ -570,9 +538,7 @@ def _scan_ollama_dir( if tmp_path.is_symlink() or tmp_path.exists(): tmp_path.unlink() except OSError as cleanup_err: - logger.debug( - "Could not clean up tmp path %s: %s", tmp_path, cleanup_err - ) + logger.debug("Could not clean up tmp path %s: %s", tmp_path, cleanup_err) return None try: @@ -589,11 +555,7 @@ def _scan_ollama_dir( repo_parts = list(parts[1:-1]) tag = parts[-1] - if ( - host == "registry.ollama.ai" - and repo_parts - and repo_parts[0] == "library" - ): + if host == "registry.ollama.ai" and repo_parts and repo_parts[0] == "library": repo_name = "/".join(repo_parts[1:]) elif host == "registry.ollama.ai": repo_name = "/".join(repo_parts) @@ -650,9 +612,7 @@ def _scan_ollama_dir( candidate = blobs_dir / digest.replace(":", "-") if candidate.is_file(): link_name = f"{safe_name}-{tag}{quant}.gguf" - gguf_link_path = _make_link( - model_link_dir, link_name, candidate - ) + gguf_link_path = _make_link(model_link_dir, link_name, candidate) elif media == "application/vnd.ollama.image.projector": candidate = blobs_dir / digest.replace(":", "-") @@ -699,10 +659,7 @@ async def list_local_models( ), current_subject: str = Depends(get_current_subject), ): - """ - List local model candidates from custom models dir, HF cache, - legacy Unsloth HF cache, and LM Studio directories. - """ + """List local model candidates from the models dir, HF caches, and LM Studio dirs.""" from utils.paths import ( legacy_hf_cache_dir, hf_default_cache_dir, @@ -715,9 +672,9 @@ async def list_local_models( hf_default = hf_default_cache_dir() lm_dirs = lmstudio_model_dirs() - # Validate models_dir against an allowlist of trusted directories. - # Only the trusted Path objects are used for filesystem access -- the - # user-supplied string is only used for matching, never for path construction. + # Validate models_dir against an allowlist of trusted dirs. Only the + # trusted Path objects are used for FS access; the user string is + # used for matching only, never for path construction. allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir] if legacy_hf.is_dir(): allowed_roots.append(legacy_hf) @@ -725,7 +682,6 @@ async def list_local_models( allowed_roots.append(hf_default) try: from utils.paths import studio_root, outputs_root - allowed_roots.extend([studio_root(), outputs_root()]) except Exception: pass @@ -735,7 +691,7 @@ async def list_local_models( for root in allowed_roots: root_str = os.path.realpath(str(root)) if requested == root_str or requested.startswith(root_str + os.sep): - models_root = root # Use the trusted root, not the user-supplied path + models_root = root # trusted root, not the user-supplied path break if models_root is None: raise HTTPException( @@ -746,11 +702,11 @@ async def list_local_models( try: local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) - # Scan legacy Unsloth HF cache for backward compatibility + # Scan legacy Unsloth HF cache for backward compatibility. if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve(): local_models += _scan_hf_cache(legacy_hf) - # Scan HF system default cache (may differ when env vars are overridden) + # Scan HF system default cache (may differ under env overrides). if ( hf_default.is_dir() and hf_default.resolve() != hf_cache_dir.resolve() @@ -758,11 +714,11 @@ async def list_local_models( ): local_models += _scan_hf_cache(hf_default) - # Scan LM Studio directories + # Scan LM Studio directories. for lm_dir in lm_dirs: local_models += _scan_lmstudio_dir(lm_dir) - # Scan user-added custom folders (cap per-folder to avoid unbounded scans) + # Scan user-added custom folders (per-folder cap). from storage.studio_db import list_scan_folders _MAX_MODELS_PER_FOLDER = 200 @@ -774,9 +730,8 @@ async def list_local_models( for folder in custom_folders: folder_path = Path(folder["path"]) try: - # Ollama scanner creates .studio_links/ with .gguf symlinks. - # Filter those from the generic scanners to avoid duplicates - # and leaking internal paths into the UI. + # Filter Ollama .studio_links/ from generic scanners to + # avoid duplicates and leaking internal paths into the UI. _generic = [ m for m in ( @@ -784,10 +739,7 @@ async def list_local_models( + _scan_hf_cache(folder_path) + _scan_lmstudio_dir(folder_path) ) - if not any( - p in (".studio_links", "ollama_links") - for p in Path(m.path).parts - ) + if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts) ] custom_models = _generic if len(custom_models) < _MAX_MODELS_PER_FOLDER: @@ -798,14 +750,11 @@ async def list_local_models( except OSError as e: logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e) continue - local_models += [ - m.model_copy(update = {"source": "custom"}) for m in custom_models - ] + local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models] - # Deduplicate models, but always keep custom folder entries so they - # appear in the "Custom Folders" UI section even when the same model - # also exists in the HF cache or default models directory. Use a - # (id, source) key for custom entries to avoid collisions. + # Deduplicate, but always keep custom folder entries (keyed by + # (id, source)) so they show in the "Custom Folders" UI section + # even when the model is also in the HF cache. deduped: dict[str, LocalModelInfo] = {} for model in local_models: key = f"{model.id}\x00custom" if model.source == "custom" else model.id @@ -817,6 +766,7 @@ async def list_local_models( key = lambda item: (item.updated_at or 0), reverse = True, ) + models = [m for m in models if not _is_hidden_model(m.id, m.path)] return LocalModelListResponse( models_dir = str(models_root), @@ -825,27 +775,25 @@ async def list_local_models( models = models, ) except Exception as e: - logger.error(f"Error listing local models: {e}", exc_info = True) - raise HTTPException( - status_code = 500, - detail = f"Failed to list local models: {str(e)}", + raise log_and_http_error( + e, + 500, + "Failed to list local models", + event = "models.list_local_models_failed", + log = logger, ) @router.get("/scan-folders") -async def get_scan_folders( - current_subject: str = Depends(get_current_subject), -): +async def get_scan_folders(current_subject: str = Depends(get_current_subject)): """List all registered custom model scan folders.""" from storage.studio_db import list_scan_folders - return {"folders": list_scan_folders()} @router.post("/scan-folders", response_model = ScanFolderInfo, status_code = 201) async def add_scan_folder_endpoint( - body: AddScanFolderRequest, - current_subject: str = Depends(get_current_subject), + body: AddScanFolderRequest, current_subject: str = Depends(get_current_subject) ): """Register a new directory to scan for local models.""" from storage.studio_db import add_scan_folder @@ -854,15 +802,16 @@ async def add_scan_folder_endpoint( folder = add_scan_folder(body.path) except ValueError as e: logger.warning("Scan folder rejected: %s (path=%s)", e, body.path) - raise HTTPException(status_code = 400, detail = str(e)) + # Forward the curated, path-free validation message. + rejection_message = str(e) + raise HTTPException(status_code = 400, detail = rejection_message) logger.info("Scan folder added: %s", folder.get("path")) return folder @router.delete("/scan-folders/{folder_id}") async def remove_scan_folder_endpoint( - folder_id: int, - current_subject: str = Depends(get_current_subject), + folder_id: int, current_subject: str = Depends(get_current_subject) ): """Remove a registered custom scan folder.""" from storage.studio_db import remove_scan_folder @@ -873,16 +822,12 @@ async def remove_scan_folder_endpoint( @router.get("/recommended-folders") -async def get_recommended_folders( - current_subject: str = Depends(get_current_subject), -): +async def get_recommended_folders(current_subject: str = Depends(get_current_subject)): """Return well-known model directories that exist on this machine. - Lightweight alternative to ``browse-folders`` for showing quick-pick - chips without the overhead of enumerating a directory tree. Returns - paths that actually exist on disk (HF cache, LM Studio, Ollama, - ``~/models``, etc.) so the frontend can offer them as one-click - "Recommended" shortcuts in the Custom Folders section. + Lightweight alternative to ``browse-folders`` for the frontend's + one-click "Recommended" chips; returns existing paths only (HF + cache, LM Studio, Ollama, ``~/models``, etc.). """ from utils.paths.storage_roots import lmstudio_model_dirs @@ -902,14 +847,14 @@ async def get_recommended_folders( seen.add(resolved) folders.append(resolved) - # LM Studio model directories + # LM Studio model directories. try: for p in lmstudio_model_dirs(): _add(p) except Exception as e: logger.warning("Failed to scan for LM Studio model directories: %s", e) - # Ollama model directories + # Ollama model directories. ollama_env = os.environ.get("OLLAMA_MODELS") if ollama_env: _add(Path(ollama_env).expanduser()) @@ -923,26 +868,20 @@ async def get_recommended_folders( return {"folders": folders} -# Heuristic ceiling on how many children to stat when checking whether a -# directory "looks like" it contains models. Keeps the browser snappy -# even when a directory has thousands of unrelated entries. +# Max children to stat when checking if a directory "looks like" it +# holds models; keeps the browser snappy on huge dirs. _BROWSE_MODEL_HINT_PROBE = 64 -# Hard cap on how many subdirectory entries we send back. Pointing the -# browser at something like ``/usr/lib`` or ``/proc`` must not stat-storm -# the process or send tens of thousands of rows to the client. +# Hard cap on subdirectory entries returned, so browsing ``/usr/lib`` +# can't stat-storm the process or flood the client. _BROWSE_ENTRY_CAP = 2000 def _count_model_files(directory: Path, cap: int = 200) -> int: """Count GGUF/safetensors files immediately inside *directory*. - Used to surface a count-hint on the response so the UI can tell - users that a leaf directory (no subdirs, only weights) is a valid - "Use this folder" target. - Bounded by *visited entries*, not by *match count*: in directories - with many non-model files (or many subdirectories) the scan still - stops after ``cap`` entries so a UI hint never costs more than a - bounded directory walk. + Surfaces a count-hint so the UI can mark a weights-only leaf dir as a + valid "Use this folder" target. Bounded by *visited entries* (stops + after ``cap``), so the hint never costs more than a bounded walk. """ n = 0 visited = 0 @@ -968,10 +907,9 @@ def _count_model_files(directory: Path, cap: int = 200) -> int: def _has_direct_model_signal(directory: Path) -> bool: - """Return True if *directory* has an immediate child that signals - it holds a model: a GGUF/safetensors/config.json file, or a - `models--*` subdir (HF hub cache). Bounded by - ``_BROWSE_MODEL_HINT_PROBE`` to stay fast.""" + """Return True if an immediate child signals a model: a + GGUF/safetensors/config.json file or a ``models--*`` subdir (HF + cache). Bounded by ``_BROWSE_MODEL_HINT_PROBE``.""" try: it = directory.iterdir() except OSError: @@ -998,23 +936,13 @@ def _has_direct_model_signal(directory: Path) -> bool: def _looks_like_model_dir(directory: Path) -> bool: - """Bounded heuristic used by the folder browser to flag directories - worth exploring. False negatives are fine; the real scanner is - authoritative. + """Bounded heuristic to flag dirs worth exploring in the browser. - Three signals, cheapest first: - - 1. Directory name itself: ``models--*`` is the HuggingFace hub cache - layout (``blobs``/``refs``/``snapshots`` children wouldn't match - the file-level probes below). - 2. An immediate child is a weight file or config (handled by - :func:`_has_direct_model_signal`). - 3. A grandchild has a direct signal -- this catches the - ``publisher/model/weights.gguf`` layout used by LM Studio and - Ollama. We probe at most the first - ``_BROWSE_MODEL_HINT_PROBE`` child directories, each of which is - checked with a bounded :func:`_has_direct_model_signal` call, - so the total cost stays O(PROBE^2) worst-case. + False negatives are fine (the real scanner is authoritative). Three + signals, cheapest first: (1) name ``models--*`` (HF cache layout), + (2) an immediate child weight/config file, (3) a grandchild with a + direct signal (LM Studio / Ollama ``publisher/model`` layout, probing + the first ``_BROWSE_MODEL_HINT_PROBE`` child dirs). """ if directory.name.startswith("models--"): return True @@ -1034,7 +962,6 @@ def _looks_like_model_dir(directory: Path) -> bool: continue except OSError: continue - # Fast name check first if child.name.startswith("models--"): return True if _has_direct_model_signal(child): @@ -1045,16 +972,13 @@ def _looks_like_model_dir(directory: Path) -> bool: def _build_browse_allowlist() -> list[Path]: - """Return the list of root directories the folder browser is allowed - to walk. The same list is used to seed the sidebar suggestion chips, - so chip targets are always reachable. + """Return the root directories the folder browser may walk. - Roots include the current user's HOME, the resolved HF cache dirs, - Studio's own outputs/exports/studio root, registered scan folders, - and well-known third-party local-LLM dirs (LM Studio, Ollama, - `~/models`). Each is added only if it currently resolves to a real - directory, so we never produce a "dead" sandbox boundary the user - can't navigate into. + The same list seeds the sidebar suggestion chips, so chip targets are + always reachable. Roots: HOME, resolved HF cache dirs, Studio's + outputs/exports/studio root, registered scan folders, and well-known + local-LLM dirs (LM Studio, Ollama, ``~/models``); each added only if + it resolves to a real directory. """ from utils.paths import ( hf_default_cache_dir, @@ -1123,9 +1047,9 @@ def _build_browse_allowlist() -> list[Path]: def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool: - """Return True if *target* equals or is a descendant of any allowed - root. The comparison uses ``os.path.realpath`` so symlinks cannot be - used to escape the sandbox. + """True if *target* equals or descends from any allowed root. + + Uses ``os.path.realpath`` so symlinks can't escape the sandbox. """ try: target_real = os.path.realpath(str(target)) @@ -1182,12 +1106,13 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]: except PermissionError: raise HTTPException( status_code = 403, - detail = f"Permission denied reading {current}", + detail = f"Permission denied reading {current.name}", ) from None except OSError as exc: + logger.warning("browse-folders: could not read %s: %s", current, exc, exc_info = True) raise HTTPException( status_code = 500, - detail = f"Could not read {current}: {exc}", + detail = f"Could not read {os.path.basename(str(current))}", ) from exc return None @@ -1219,14 +1144,21 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa if child is None: raise HTTPException( status_code = 404, - detail = f"Path does not exist: {requested_path}", + detail = f"Path does not exist: {os.path.basename(requested_path)}", ) try: resolved_child = child.resolve() except OSError as exc: + logger.warning( + "browse-folders: invalid path component %r under %s: %s", + part, + current, + exc, + exc_info = True, + ) raise HTTPException( status_code = 400, - detail = f"Invalid path: {exc}", + detail = "Invalid path", ) from exc if not _is_path_inside_allowlist(resolved_child, resolved_roots): raise HTTPException( @@ -1242,7 +1174,7 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa if not current.is_dir(): raise HTTPException( status_code = 400, - detail = f"Not a directory: {current}", + detail = f"Not a directory: {os.path.basename(str(current))}", ) return current @@ -1274,33 +1206,21 @@ async def browse_folders( ), current_subject: str = Depends(get_current_subject), ): - """ - List immediate subdirectories of *path* for the Custom Folders picker. + """List immediate subdirectories of *path* for the Custom Folders picker. - The frontend uses this to render a modal folder browser without needing - a native OS dialog (Studio is served over HTTP, so the browser can't - reveal absolute paths on the host). The endpoint is read-only and does - not create, move, or delete anything. It simply enumerates visible - subdirectories so the user can click their way to a folder and hand - the resulting string back to POST `/api/models/scan-folders`. + Lets the frontend render a modal folder browser without a native OS + dialog. Read-only: enumerates visible subdirectories so the user can + click to a folder and hand the string to POST /api/models/scan-folders. - Sandbox: requests are bounded to the allowlist returned by - :func:`_build_browse_allowlist` (HOME, HF cache, Studio dirs, - registered scan folders, well-known model dirs). Paths outside the - allowlist return 403 so users cannot probe ``/etc``, ``/proc``, - ``/root`` (when not HOME), or other sensitive system locations - even if the server process can read them. Symlinks are resolved - via ``os.path.realpath`` before the check, so symlink traversal - cannot escape the sandbox either. - - Sorting: directories that look like they hold models come first, then - plain directories, then hidden entries (if `show_hidden=true`). + Sandbox: bounded to :func:`_build_browse_allowlist`; paths outside it + return 403, and symlinks are resolved via ``os.path.realpath`` first + so traversal can't escape. Sorting: model-bearing dirs, then plain, + then hidden (if ``show_hidden=true``). """ from utils.paths import hf_default_cache_dir, well_known_model_dirs from storage.studio_db import list_scan_folders - # Build the allowlist once -- both the sandbox check below and the - # suggestion chips use the same set, so chips are always navigable. + # Build once; the sandbox check and suggestion chips share it. allowed_roots = _build_browse_allowlist() try: @@ -1315,8 +1235,7 @@ async def browse_folders( ) raise - # Enumerate immediate subdirectories with a bounded cap so a stray - # query against ``/usr/lib`` or ``/proc`` can't stat-storm the process. + # Enumerate immediate subdirectories with a bounded cap. entries: list[BrowseEntry] = [] truncated = False visited = 0 @@ -1325,23 +1244,20 @@ async def browse_folders( except PermissionError: raise HTTPException( status_code = 403, - detail = f"Permission denied reading {target}", + detail = f"Permission denied reading {os.path.basename(str(target))}", ) except OSError as exc: + logger.warning("browse-folders: could not read %s: %s", target, exc, exc_info = True) raise HTTPException( status_code = 500, - detail = f"Could not read {target}: {exc}", + detail = f"Could not read {os.path.basename(str(target))}", ) try: for child in it: - # Bound by *visited entries*, not by *appended entries*: in - # directories full of files (or hidden subdirs when - # ``show_hidden=False``) the cap on ``len(entries)`` would - # never trigger and we'd still stat every child. Counting - # visits keeps the worst-case work to ``_BROWSE_ENTRY_CAP`` - # iterdir/is_dir calls regardless of how many of them - # survive the filters below. + # Bound by *visited*, not *appended*: a cap on len(entries) + # would never trigger in dirs full of files. Counting visits + # caps worst-case work at ``_BROWSE_ENTRY_CAP`` calls. visited += 1 if visited > _BROWSE_ENTRY_CAP: truncated = True @@ -1369,10 +1285,10 @@ async def browse_folders( exc, ) except OSError as exc: - # Rare: iterdir succeeded but reading a specific entry failed. + # Rare: iterdir succeeded but reading an entry failed. logger.warning("browse-folders: partial enumeration of %s: %s", target, exc) - # Model-bearing dirs first, then plain, then hidden; case-insensitive + # Model-bearing first, then plain, then hidden; case-insensitive # alphabetical within each bucket. def _sort_key(e: BrowseEntry) -> tuple[int, str]: bucket = 0 if e.has_models else (2 if e.hidden else 1) @@ -1380,14 +1296,11 @@ async def browse_folders( entries.sort(key = _sort_key) - # Parent is None at the filesystem root (`p.parent == p`) AND when - # the parent would step outside the sandbox -- otherwise the up-row - # would 403 on click. Users can still hop to other allowed roots - # via the suggestion chips below. + # Parent is None at the filesystem root and when it would leave the + # sandbox (else the up-row would 403 on click); users can still hop + # to other allowed roots via the suggestion chips. parent: Optional[str] - if target.parent == target or not _is_path_inside_allowlist( - target.parent, allowed_roots - ): + if target.parent == target or not _is_path_inside_allowlist(target.parent, allowed_roots): parent = None else: parent = str(target.parent) @@ -1409,27 +1322,21 @@ async def browse_folders( seen_sug.add(resolved) suggestions.append(resolved) - # Home always comes first -- it's the safe fallback when everything - # else is cold. + # Home first -- the safe fallback when everything else is cold. _add_sug(Path.home()) # The HF cache root the process is actually using. try: _add_sug(hf_default_cache_dir()) except Exception: pass - # Already-registered scan folders (what the user has curated). + # Already-registered scan folders (user-curated). try: for folder in list_scan_folders(): _add_sug(Path(folder.get("path", ""))) except Exception as exc: logger.debug("browse-folders: could not load scan folders: %s", exc) - # Directories commonly used by other local-LLM tools: LM Studio - # (`~/.lmstudio/models` + legacy `~/.cache/lm-studio/models` + - # user-configured downloadsFolder from LM Studio's settings.json), - # Ollama (`~/.ollama/models` + common system paths + OLLAMA_MODELS - # env var), and generic user-choice spots (`~/models`, `~/Models`). - # Each helper only returns paths that currently exist so we never - # show dead chips. + # Dirs used by other local-LLM tools (LM Studio, Ollama, ~/models); + # the helper returns only existing paths, so no dead chips. try: for p in well_known_model_dirs(): _add_sug(p) @@ -1446,22 +1353,23 @@ async def browse_folders( ) -@router.get("/list") -async def list_models( - current_subject: str = Depends(get_current_subject), -): - """ - List available models (default models and loaded models). +def _looks_like_mlx_repo(model_id: str) -> bool: + """Name heuristic for unloaded models (mirrors the -GGUF suffix check); + tokenized so MLX only matches as a whole name segment.""" + if model_id.lower().startswith("mlx-community/"): + return True + tail = model_id.split("/")[-1] + return "MLX" in _re.split(r"[-_.]", tail.upper()) - This endpoint returns the default models and any currently loaded models. - """ + +@router.get("/list") +async def list_models(current_subject: str = Depends(get_current_subject)): + """List available models: default plus currently loaded.""" try: inference_backend = get_inference_backend() - # Get default models default_models = inference_backend.default_models - # Get loaded models loaded_models = [] for model_name, model_data in inference_backend.models.items(): _is_vision = model_data.get("is_vision", False) @@ -1471,6 +1379,7 @@ async def list_models( name = model_name.split("/")[-1] if "/" in model_name else model_name, is_vision = _is_vision, is_lora = model_data.get("is_lora", False), + is_mlx = model_data.get("is_mlx", False), is_audio = model_data.get("is_audio", False), audio_type = _audio_type, has_audio_input = model_data.get("has_audio_input", False), @@ -1478,7 +1387,7 @@ async def list_models( ) loaded_models.append(model_info) - # Include active GGUF model (loaded via llama-server) + # Include active GGUF model (loaded via llama-server). from routes.inference import get_llama_cpp_backend llama_backend = get_llama_cpp_backend() @@ -1494,22 +1403,23 @@ async def list_models( ) ) - # Combine default and loaded models + # Combine default and loaded; prefer loaded entries for duplicate + # ids so runtime flags survive. all_models = [] seen_ids = set() + loaded_by_id = {model_info.id: model_info for model_info in loaded_models} - # Add default models for model_id in default_models: if model_id not in seen_ids: - model_info = ModelDetails( + model_info = loaded_by_id.get(model_id) or ModelDetails( id = model_id, name = model_id.split("/")[-1] if "/" in model_id else model_id, is_gguf = model_id.upper().endswith("-GGUF"), + is_mlx = _looks_like_mlx_repo(model_id), ) all_models.append(model_info) seen_ids.add(model_id) - # Add loaded models for model_info in loaded_models: if model_info.id not in seen_ids: all_models.append(model_info) @@ -1518,25 +1428,26 @@ async def list_models( return ModelListResponse(models = all_models, default_models = default_models) except Exception as e: - logger.error(f"Error listing models: {e}", exc_info = True) - raise HTTPException(status_code = 500, detail = f"Failed to list models: {str(e)}") + raise log_and_http_error( + e, + 500, + "Failed to list models", + event = "models.list_models_failed", + log = logger, + ) def _get_max_position_embeddings(config) -> Optional[int]: - """Extract max_position_embeddings from a model config, checking text_config fallback.""" + """Extract max_position_embeddings from a config, with text_config fallback.""" if hasattr(config, "max_position_embeddings"): return config.max_position_embeddings - if hasattr(config, "text_config") and hasattr( - config.text_config, "max_position_embeddings" - ): + if hasattr(config, "text_config") and hasattr(config.text_config, "max_position_embeddings"): return config.text_config.max_position_embeddings return None -def _get_model_size_bytes( - model_name: str, hf_token: Optional[str] = None -) -> Optional[int]: - """Get total size of model weight files from HF Hub.""" +def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Optional[int]: + """Total size of model weight files from HF Hub.""" try: from huggingface_hub import HfApi @@ -1548,9 +1459,7 @@ def _get_model_size_bytes( weight_exts = (".safetensors", ".bin", ".pt", ".pth", ".gguf") total = 0 for sibling in info.siblings: - if sibling.rfilename and any( - sibling.rfilename.endswith(ext) for ext in weight_exts - ): + if sibling.rfilename and any(sibling.rfilename.endswith(ext) for ext in weight_exts): if sibling.size is not None: total += sibling.size @@ -1566,11 +1475,7 @@ async def get_model_config( hf_token: Optional[str] = Query(None), current_subject: str = Depends(get_current_subject), ): - """ - Get configuration for a specific model. - - This endpoint wraps the backend load_model_defaults function. - """ + """Get configuration for a specific model (wraps load_model_defaults).""" try: if not is_local_path(model_name): resolved = resolve_cached_repo_id_case(model_name) @@ -1585,15 +1490,13 @@ async def get_model_config( logger.info(f"Getting model config for: {model_name}") from utils.models.model_config import detect_audio_type - # Load model defaults from backend config_dict = load_model_defaults(model_name) - # Detect model capabilities (pass HF token for gated models) + # Detect capabilities (HF token for gated models). is_vision = is_vision_model(model_name, hf_token = hf_token) is_embedding = is_embedding_model(model_name, hf_token = hf_token) audio_type = detect_audio_type(model_name, hf_token = hf_token) - # Check if it's a LoRA adapter is_lora = False base_model = None max_position_embeddings = None @@ -1605,7 +1508,7 @@ async def get_model_config( except Exception: pass - # Fallback: try AutoConfig directly if not found yet + # Fallback: try AutoConfig directly. if max_position_embeddings is None: try: from transformers import AutoConfig as _AutoConfig @@ -1638,9 +1541,12 @@ async def get_model_config( ) except Exception as e: - logger.error(f"Error getting model config: {e}", exc_info = True) - raise HTTPException( - status_code = 500, detail = f"Failed to get model config: {str(e)}" + raise log_and_http_error( + e, + 500, + "Failed to get model config", + event = "models.get_model_config_failed", + log = logger, ) @@ -1654,18 +1560,16 @@ async def scan_loras( ), current_subject: str = Depends(get_current_subject), ): - """ - Scan for trained LoRA adapters and exported models. + """Scan for trained LoRA adapters and exported models. - Returns both training outputs (from outputs_dir) and exported models - (from exports_dir) in a single list, distinguished by source field. + Returns training outputs (outputs_dir) and exported models + (exports_dir) in one list, distinguished by the source field. """ try: resolved_outputs_dir = str(resolve_output_dir(outputs_dir)) resolved_exports_dir = str(resolve_export_dir(exports_dir)) lora_list = [] - # Scan training outputs trained_models = scan_trained_models(outputs_dir = resolved_outputs_dir) for display_name, model_path, model_type in trained_models: base_model = get_base_model_from_checkpoint(model_path) @@ -1695,9 +1599,12 @@ async def scan_loras( return LoRAScanResponse(loras = lora_list, outputs_dir = resolved_outputs_dir) except Exception as e: - logger.error(f"Error scanning LoRAs: {e}", exc_info = True) - raise HTTPException( - status_code = 500, detail = f"Failed to scan LoRA adapters: {str(e)}" + raise log_and_http_error( + e, + 500, + "Failed to scan LoRA adapters", + event = "models.scan_loras_failed", + log = logger, ) @@ -1710,7 +1617,7 @@ def _is_path_under(path: Path, root: Path) -> bool: def _is_path_under_lexically(path: Path, root: Path) -> bool: - """Check containment without resolving the final path's symlink target.""" + """Check containment without resolving the final path's symlink.""" try: absolute_path = Path(os.path.abspath(str(path))) absolute_root = Path(os.path.abspath(str(root))) @@ -1732,25 +1639,20 @@ def _loaded_model_matches_deleted_path(active_model: str, deleted_path: Path) -> ) active_lower = active_model.lower() target_lower = str(deleted_path).lower() - return active_lower == target_lower or active_lower.startswith( - f"{target_lower}{os.sep}" - ) + return active_lower == target_lower or active_lower.startswith(f"{target_lower}{os.sep}") -def _loading_model_matches_deleted_path( - loading_model: object, - deleted_path: Path, -) -> bool: +def _loading_model_matches_deleted_path(loading_model: object, deleted_path: Path) -> bool: if not loading_model: return False return _loaded_model_matches_deleted_path(str(loading_model), deleted_path) def _prune_empty_parents(start: Path, stop_at: Path) -> None: - """Remove empty ancestor directories of ``start`` up to (but not including) ``stop_at``. + """Remove empty ancestors of ``start`` up to (not including) ``stop_at``. - Used after deleting a model checkpoint so the enclosing run directory does - not linger as an empty entry in scan results. + Used after deleting a checkpoint so the enclosing run dir doesn't + linger as an empty entry in scan results. """ try: stop_resolved = stop_at.resolve() @@ -1802,8 +1704,8 @@ async def delete_finetuned_model( ): """Delete a Studio-trained or exported model from disk. - Only paths under Studio's outputs/exports roots are accepted. Exported - GGUF entries can delete one quantization variant at a time. + Only paths under Studio's outputs/exports roots are accepted. + Exported GGUF entries can delete one quant variant at a time. """ if source not in {"training", "exported"}: raise HTTPException( @@ -1873,7 +1775,6 @@ async def delete_finetuned_model( if source == "training": try: from core.training import get_training_backend - training_backend = get_training_backend() if training_backend.is_training_active(): raise HTTPException( @@ -1960,9 +1861,7 @@ async def delete_finetuned_model( except HTTPException: raise except Exception as e: - logger.warning( - "Could not check inference backend loaded model before delete: %s", e - ) + logger.warning("Could not check inference backend loaded model before delete: %s", e) raise HTTPException( status_code = 503, detail = "Could not verify model load status before deleting", @@ -2029,15 +1928,12 @@ async def delete_finetuned_model( ) raise HTTPException( status_code = 500, - detail = f"Failed to delete fine-tuned model: {str(e)}", + detail = "Failed to delete fine-tuned model", ) @router.get("/loras/{lora_path:path}/base-model", response_model = LoRABaseModelResponse) -async def get_lora_base_model( - lora_path: str, - current_subject: str = Depends(get_current_subject), -): +async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get_current_subject)): """ Get the base model for a LoRA adapter. @@ -2060,17 +1956,17 @@ async def get_lora_base_model( except HTTPException: raise except Exception as e: - logger.error(f"Error getting LoRA base model: {e}", exc_info = True) - raise HTTPException( - status_code = 500, detail = f"Failed to get base model: {str(e)}" + raise log_and_http_error( + e, + 500, + "Failed to get base model", + event = "models.get_lora_base_model_failed", + log = logger, ) @router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse) -async def check_vision_model( - model_name: str, - current_subject: str = Depends(get_current_subject), -): +async def check_vision_model(model_name: str, current_subject: str = Depends(get_current_subject)): """ Check if a model is a vision model. @@ -2087,9 +1983,12 @@ async def check_vision_model( ) except Exception as e: - logger.error(f"Error checking vision model: {e}", exc_info = True) - raise HTTPException( - status_code = 500, detail = f"Failed to check vision model: {str(e)}" + raise log_and_http_error( + e, + 500, + "Failed to check vision model", + event = "models.check_vision_model_failed", + log = logger, ) @@ -2108,18 +2007,19 @@ async def check_embedding_model( logger.info(f"Checking if embedding model: {model_name}") is_embedding = is_embedding_model(model_name, hf_token = hf_token) - logger.info( - f"Embedding check result for {model_name}: is_embedding={is_embedding}" - ) + logger.info(f"Embedding check result for {model_name}: is_embedding={is_embedding}") return EmbeddingCheckResponse( model_name = model_name, is_embedding = is_embedding, ) except Exception as e: - logger.error(f"Error checking embedding model: {e}", exc_info = True) - raise HTTPException( - status_code = 500, detail = f"Failed to check embedding model: {str(e)}" + raise log_and_http_error( + e, + 500, + "Failed to check embedding model", + event = "models.check_embedding_model_failed", + log = logger, ) @@ -2128,23 +2028,18 @@ async def get_gguf_variants( repo_id: str = Query( ..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')" ), - hf_token: Optional[str] = Query( - None, description = "HuggingFace token for private repos" - ), + hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"), current_subject: str = Depends(get_current_subject), ): - """ - List available GGUF quantization variants for a HuggingFace repo - or a local directory (e.g. LM Studio model folder). + """List GGUF quantization variants for a HF repo or local directory. - Returns all available quantization variants (Q4_K_M, Q8_0, BF16, etc.) - with file sizes, whether the model supports vision, and the recommended - default variant. + Returns all variants with file sizes, vision support, and the + recommended default. """ try: from utils.models.model_config import is_local_path, list_local_gguf_variants - # Local directory path (e.g. LM Studio models) — scan filesystem + # Local directory path — scan filesystem. if is_local_path(repo_id): variants, has_vision = list_local_gguf_variants(repo_id) @@ -2167,26 +2062,22 @@ async def get_gguf_variants( default_variant = default_variant, ) - # Remote HuggingFace repo — query HF API + # Remote HuggingFace repo — query HF API. variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token) - # Determine default variant filenames = [v.filename for v in variants] best = _pick_best_gguf(filenames) default_variant = _extract_quant_label(best) if best else None - # Check which variants are fully downloaded in the HF cache. - # For split GGUFs, ALL shards must be present -- sum cached bytes - # per variant and compare against the expected total. - # HF cache dir uses the exact case from the repo_id at download time, - # which may differ from the canonical HF repo_id, so do a - # case-insensitive match. + # Which variants are fully downloaded in the HF cache. For split + # GGUFs ALL shards must be present, so sum cached bytes per variant + # vs. the expected total. Cache dir casing may differ from the + # canonical repo_id, so match case-insensitively. cached_bytes_by_quant: dict[str, int] = {} try: import re as _re from huggingface_hub import constants as hf_constants - # Sanitize repo_id: must be "owner/name" with safe chars only if not _is_valid_repo_id(repo_id): raise ValueError(f"Invalid repo_id format: {repo_id}") @@ -2210,7 +2101,7 @@ async def get_gguf_variants( cached = cached_bytes_by_quant.get(variant.quant, 0) if cached == 0 or variant.size_bytes == 0: return False - # Allow small rounding tolerance (symlinks vs real sizes) + # Rounding tolerance (symlinks vs real sizes). return cached >= variant.size_bytes * 0.99 return GgufVariantsResponse( @@ -2232,7 +2123,7 @@ async def get_gguf_variants( logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info = True) raise HTTPException( status_code = 500, - detail = f"Failed to list GGUF variants: {str(e)}", + detail = "Failed to list GGUF variants", ) @@ -2243,10 +2134,10 @@ async def get_gguf_download_progress( expected_bytes: int = Query(0, description = "Expected total download size in bytes"), current_subject: str = Depends(get_current_subject), ): - """Return download progress by checking cached GGUF files for a specific variant. + """Download progress from cached GGUF files for a specific variant. - Tracks completed shard downloads in snapshots and in-progress downloads - in the blobs directory (incomplete files). + Tracks completed shards in snapshots and in-progress (.incomplete) + downloads in the blobs directory. """ try: if not _is_valid_repo_id(repo_id): @@ -2265,12 +2156,12 @@ async def get_gguf_download_progress( in_progress_bytes = 0 for entry in cache_dir.iterdir(): if entry.name.lower() == target: - # Count completed .gguf files matching this variant in snapshots + # Completed .gguf files for this variant in snapshots. for f in _iter_gguf_paths(entry): fname = f.name.lower().replace("-", "").replace("_", "") if not variant_lower or variant_lower in fname: downloaded_bytes += f.stat().st_size - # Check blobs for in-progress downloads (.incomplete files) + # In-progress (.incomplete) downloads in blobs. blobs_dir = entry / "blobs" if blobs_dir.is_dir(): for f in blobs_dir.iterdir(): @@ -2279,12 +2170,8 @@ async def get_gguf_download_progress( break total_progress_bytes = downloaded_bytes + in_progress_bytes - progress = ( - min(total_progress_bytes / expected_bytes, 0.99) - if expected_bytes > 0 - else 0 - ) - # Only report 1.0 when all bytes are in completed files (not in-progress) + progress = min(total_progress_bytes / expected_bytes, 0.99) if expected_bytes > 0 else 0 + # Report 1.0 only when all bytes are in completed files. if expected_bytes > 0 and downloaded_bytes >= expected_bytes: progress = 1.0 return { @@ -2299,9 +2186,9 @@ async def get_gguf_download_progress( def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]: """Pick the most useful on-disk path for a HF cache repo. - Prefers the most-recent snapshot dir (what `from_pretrained` actually - points at). Falls back to the cache repo root. Returns the resolved - realpath so symlinks under snapshots/ are followed back to blobs/. + Prefers the most-recent snapshot dir (what ``from_pretrained`` uses), + falling back to the cache repo root. Returns the resolved realpath so + snapshot symlinks follow back to blobs/. """ try: snapshots_dir = repo_dir / "snapshots" @@ -2323,11 +2210,10 @@ async def get_download_progress( """Return download progress for any HuggingFace model repo. Checks the local HF cache for completed blobs and in-progress - (.incomplete) downloads. Uses the HF API to determine the expected - total size on the first call, then caches it for subsequent polls. - Also returns ``cache_path``: the realpath of the snapshot directory - (or the cache repo root if no snapshot exists yet) so the UI can - show users where the weights actually live on disk. + (.incomplete) downloads. Gets the expected total size from the HF API + on the first call, then caches it for later polls. Also returns + ``cache_path``: the realpath of the snapshot dir (or cache repo root + if no snapshot yet) so the UI can show where weights live on disk. """ _empty = { "downloaded_bytes": 0, @@ -2367,10 +2253,9 @@ async def get_download_progress( if downloaded_bytes == 0: return {**_empty, "cache_path": cache_path} - # Get expected size from HF API (cached per repo_id) expected_bytes = _get_repo_size_cached(repo_id) if expected_bytes <= 0: - # Cannot determine total; report bytes only, no percentage + # Total unknown; report bytes only, no percentage. return { "downloaded_bytes": downloaded_bytes, "expected_bytes": 0, @@ -2378,11 +2263,9 @@ async def get_download_progress( "cache_path": cache_path, } - # Use 95% threshold for completion (blob deduplication can make - # completed_bytes differ slightly from expected_bytes). - # Do NOT use "no .incomplete files" as a completion signal -- - # HF downloads files sequentially, so between files there are - # no .incomplete files even though the download is far from done. + # 95% threshold (blob dedup can skew completed_bytes). Do NOT + # treat "no .incomplete files" as done: HF downloads sequentially, + # so none exist between files even when far from finished. if completed_bytes >= expected_bytes * 0.95: progress = 1.0 else: @@ -2417,16 +2300,15 @@ def _get_repo_size_cached(repo_id: str) -> int: def _all_hf_cache_scans(): - """Return scan_cache_dir results for the active, legacy, and default HF caches.""" + """scan_cache_dir results for the active, legacy, and default HF caches.""" from huggingface_hub import scan_cache_dir from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir scans = [scan_cache_dir()] seen: set[str] = set() try: - # Resolve the active cache dir so we can dedup + # Resolve the active cache dir for dedup. from huggingface_hub.constants import HF_HUB_CACHE - seen.add(str(Path(HF_HUB_CACHE).resolve())) except Exception: pass @@ -2447,14 +2329,14 @@ def _is_gguf_filename(name: str) -> bool: def _is_mmproj_filename(name: str) -> bool: - """Match GGUF vision-adapter (mmproj) files. Kept consistent with + """Match GGUF vision-adapter (mmproj) files. Consistent with ``utils.models.model_config._is_mmproj``.""" return "mmproj" in name.lower() def _is_main_gguf_filename(name: str) -> bool: - """A GGUF file that is a primary weight artifact, not an mmproj - vision adapter.""" + """A GGUF file that is a primary weight, not an mmproj vision + adapter.""" return _is_gguf_filename(name) and not _is_mmproj_filename(name) @@ -2465,17 +2347,16 @@ def _iter_gguf_paths(root: Path): def _repo_gguf_size_bytes(repo_info) -> int: - """Return the total on-disk size of primary GGUF weight files across - all revisions, excluding mmproj vision-adapter files. + """Total on-disk size of primary GGUF weight files across all + revisions, excluding mmproj vision-adapter files. Hugging Face hardlinks blobs shared between revisions, so this - deduplicates by blob path (or, as a fallback, by revision commit - hash + filename) to avoid double-counting the same bytes. Files - with an unknown size (``size_on_disk is None``, e.g. a partial or - interrupted download) are treated as zero bytes. mmproj files are - excluded so that repos whose only ``.gguf`` artifact is a vision - adapter are not classified as GGUF repos: the variant selector - filters mmproj out and would otherwise show zero pickable variants. + deduplicates by blob path (or revision commit hash + filename as a + fallback) to avoid double-counting. Unknown sizes (``size_on_disk is + None``, e.g. a partial download) count as zero. mmproj files are + excluded so repos whose only ``.gguf`` artifact is a vision adapter + aren't classed as GGUF repos: the variant selector filters mmproj + out and would otherwise show zero pickable variants. """ unique_blobs: dict[str, int] = {} for revision in repo_info.revisions: @@ -2492,16 +2373,14 @@ def _repo_gguf_size_bytes(repo_info) -> int: def _repo_has_gguf_files(repo_info) -> bool: - """Return True when any revision in a cached repo contains a - primary GGUF weight file. Repos whose only ``.gguf`` artifact is - an mmproj vision adapter are not treated as GGUF here.""" + """True when any revision in a cached repo has a primary GGUF weight + file. Repos whose only ``.gguf`` artifact is an mmproj vision adapter + are not treated as GGUF here.""" return _repo_gguf_size_bytes(repo_info) > 0 @router.get("/cached-gguf") -async def list_cached_gguf( - current_subject: str = Depends(get_current_subject), -): +async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): """List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" try: cache_scans = _all_hf_cache_scans() @@ -2513,6 +2392,8 @@ async def list_cached_gguf( if repo_info.repo_type != "model": continue repo_id = repo_info.repo_id + if _is_hidden_model(repo_id): + continue total_size = _repo_gguf_size_bytes(repo_info) if total_size == 0: continue @@ -2536,9 +2417,7 @@ async def list_cached_gguf( @router.get("/cached-models") -async def list_cached_models( - current_subject: str = Depends(get_current_subject), -): +async def list_cached_models(current_subject: str = Depends(get_current_subject)): """List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" _WEIGHT_EXTENSIONS = (".safetensors", ".bin") @@ -2552,12 +2431,12 @@ async def list_cached_models( if repo_info.repo_type != "model": continue repo_id = repo_info.repo_id + if _is_hidden_model(repo_id): + continue if _repo_has_gguf_files(repo_info): continue total_size = sum( - (f.size_on_disk or 0) - for rev in repo_info.revisions - for f in rev.files + (f.size_on_disk or 0) for rev in repo_info.revisions for f in rev.files ) if total_size == 0: continue @@ -2594,17 +2473,16 @@ async def delete_cached_model( ): """Delete a cached model repo (or a specific GGUF variant) from the HF cache. - When *variant* is provided, only the GGUF files matching that quant label - are removed (e.g. ``UD-Q4_K_XL``). Otherwise the entire repo is deleted. - Refuses if the model is currently loaded for inference. + With *variant*, only GGUF files matching that quant label are removed + (e.g. ``UD-Q4_K_XL``); otherwise the whole repo is deleted. Refuses + if the model is currently loaded for inference. """ if not _is_valid_repo_id(repo_id): raise HTTPException(status_code = 400, detail = "Invalid repo_id format") - # Check if model is currently loaded + # Refuse if the model is currently loaded. try: from routes.inference import get_llama_cpp_backend - llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded and llama_backend.model_identifier: loaded_id = llama_backend.model_identifier.lower() @@ -2660,7 +2538,7 @@ async def delete_cached_model( quant = _extract_quant_label(f.file_name) if quant.lower() != variant.lower(): continue - # Delete the blob (actual data) and the snapshot symlink + # Delete the blob (data) and the snapshot symlink. try: blob = Path(f.blob_path) snap = Path(f.file_path) @@ -2707,7 +2585,7 @@ async def delete_cached_model( logger.error(f"Error deleting cached model {repo_id}: {e}", exc_info = True) raise HTTPException( status_code = 500, - detail = f"Failed to delete cached model: {str(e)}", + detail = "Failed to delete cached model", ) @@ -2719,8 +2597,7 @@ async def list_checkpoints( ), current_subject: str = Depends(get_current_subject), ): - """ - List available checkpoints in the outputs directory. + """List checkpoints in the outputs directory. Scans the outputs folder for training runs and their checkpoints. """ @@ -2748,8 +2625,10 @@ async def list_checkpoints( models = models, ) except Exception as e: - logger.error(f"Error listing checkpoints: {e}", exc_info = True) - raise HTTPException( - status_code = 500, - detail = f"Failed to list checkpoints: {str(e)}", + raise log_and_http_error( + e, + 500, + "Failed to list checkpoints", + event = "models.list_checkpoints_failed", + log = logger, ) diff --git a/studio/backend/routes/prompts.py b/studio/backend/routes/prompts.py new file mode 100644 index 0000000000..df81008766 --- /dev/null +++ b/studio/backend/routes/prompts.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Prompt storage API routes backed by studio.db. +""" + +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel, Field + +from auth.authentication import get_current_subject +from storage.studio_db import ( + bulk_upsert_prompt_entries, + bulk_upsert_prompt_lists, + delete_prompt_entry, + delete_prompt_list_db, + list_prompt_entries, + list_prompt_lists_db, + upsert_prompt_entry, + upsert_prompt_list, +) + +router = APIRouter() + + +class PromptEntry(BaseModel): + id: str = Field(max_length = 128) + name: str = Field(max_length = 500) + text: str = Field(max_length = 100_000) + createdAt: int + updatedAt: int + + +class PromptList(BaseModel): + id: str = Field(max_length = 128) + name: str = Field(max_length = 500) + items: list[str] = Field(max_length = 10_000) + createdAt: int + updatedAt: int + + +class BulkEntriesRequest(BaseModel): + entries: list[PromptEntry] + + +class BulkListsRequest(BaseModel): + lists: list[PromptList] + + +@router.get("/entries") +def get_entries(current_subject: str = Depends(get_current_subject)): + return {"entries": list_prompt_entries()} + + +@router.put("/entries/{entry_id}") +def put_entry( + entry_id: str, + entry: PromptEntry, + current_subject: str = Depends(get_current_subject), +): + if entry.id != entry_id: + raise HTTPException(status_code = 400, detail = "ID mismatch") + return upsert_prompt_entry(entry.model_dump()) + + +@router.delete("/entries/{entry_id}", status_code = 204) +def remove_entry(entry_id: str, current_subject: str = Depends(get_current_subject)): + delete_prompt_entry(entry_id) + + +@router.post("/entries/bulk") +def bulk_entries(req: BulkEntriesRequest, current_subject: str = Depends(get_current_subject)): + count = bulk_upsert_prompt_entries([e.model_dump() for e in req.entries]) + return {"count": count} + + +@router.get("/lists") +def get_lists(current_subject: str = Depends(get_current_subject)): + return {"lists": list_prompt_lists_db()} + + +@router.put("/lists/{list_id}") +def put_list( + list_id: str, + lst: PromptList, + current_subject: str = Depends(get_current_subject), +): + if lst.id != list_id: + raise HTTPException(status_code = 400, detail = "ID mismatch") + return upsert_prompt_list(lst.model_dump()) + + +@router.delete("/lists/{list_id}", status_code = 204) +def remove_list(list_id: str, current_subject: str = Depends(get_current_subject)): + delete_prompt_list_db(list_id) + + +@router.post("/lists/bulk") +def bulk_lists(req: BulkListsRequest, current_subject: str = Depends(get_current_subject)): + count = bulk_upsert_prompt_lists([l.model_dump() for l in req.lists]) + return {"count": count} diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 2bb1de5366..64d1c3eab3 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -4,12 +4,12 @@ """ API routes for external LLM provider management. -Provides endpoints for: - - Discovering available provider types (registry) +Endpoints: + - Discover available provider types (registry) - CRUD for saved provider configurations (no API keys stored) - - Fetching the RSA public key for API key encryption - - Testing provider connectivity - - Listing models from a provider + - Fetch the RSA public key for API key encryption + - Test provider connectivity + - List models from a provider """ import uuid @@ -40,6 +40,7 @@ from models.providers import ( ProviderUpdate, ) from storage import providers_db +from utils.utils import safe_curated_detail, log_and_http_error logger = structlog.get_logger(__name__) @@ -50,16 +51,11 @@ router = APIRouter() @router.get("/public-key") -async def get_public_key( - current_subject: str = Depends(get_current_subject), -): +async def get_public_key(current_subject: str = Depends(get_current_subject)): """Return the RSA public key PEM for client-side API key encryption. - The ``fingerprint`` field is a short SHA256 of the PEM and is meant - purely for diagnostics — a mismatch between what the frontend - captured at encrypt time and what the server reports here is a - clear signal that the keypair rotated mid-flight (e.g. the server - re-ran ``init_key_pair`` for any reason). + ``fingerprint`` is a short SHA256 of the PEM; a mismatch with what the + frontend captured at encrypt time signals the keypair rotated mid-flight. """ return { "public_key": get_public_key_pem(), @@ -71,9 +67,7 @@ async def get_public_key( @router.get("/registry", response_model = list[ProviderRegistryEntry]) -async def list_registry( - current_subject: str = Depends(get_current_subject), -): +async def list_registry(current_subject: str = Depends(get_current_subject)): """List all supported provider types with their default configurations.""" return list_available_providers() @@ -82,13 +76,9 @@ async def list_registry( @router.get("/pricing") -async def get_pricing_snapshot( - current_subject: str = Depends(get_current_subject), -): - """Static per-MTok pricing table the frontend uses to convert - upstream usage chunks into a per-turn USD cost. See - ``core/inference/pricing.py`` for sourcing notes; values reflect - the published prices as of the file's last update.""" +async def get_pricing_snapshot(current_subject: str = Depends(get_current_subject)): + """Static per-MTok pricing table the frontend uses to convert upstream + usage into per-turn USD cost. See ``core/inference/pricing.py`` for sourcing.""" return pricing_snapshot() @@ -96,9 +86,7 @@ async def get_pricing_snapshot( @router.get("/", response_model = list[ProviderResponse]) -async def list_provider_configs( - current_subject: str = Depends(get_current_subject), -): +async def list_provider_configs(current_subject: str = Depends(get_current_subject)): """List all saved provider configurations.""" rows = providers_db.list_providers() return [ @@ -117,8 +105,7 @@ async def list_provider_configs( @router.post("/", response_model = ProviderResponse, status_code = 201) async def create_provider_config( - payload: ProviderCreate, - current_subject: str = Depends(get_current_subject), + payload: ProviderCreate, current_subject: str = Depends(get_current_subject) ): """Create a new saved provider configuration (no API key stored).""" info = get_provider_info(payload.provider_type) @@ -185,8 +172,7 @@ async def update_provider_config( @router.delete("/{provider_id}", status_code = 204) async def delete_provider_config( - provider_id: str, - current_subject: str = Depends(get_current_subject), + provider_id: str, current_subject: str = Depends(get_current_subject) ): """Delete a saved provider configuration.""" deleted = providers_db.delete_provider(provider_id) @@ -199,14 +185,13 @@ async def delete_provider_config( @router.post("/test", response_model = ProviderTestResult) async def test_provider( - payload: ProviderTestRequest, - current_subject: str = Depends(get_current_subject), + payload: ProviderTestRequest, current_subject: str = Depends(get_current_subject) ): """ Test connectivity to an external provider. Makes a lightweight GET /models call to verify the API key works. - The encrypted_api_key is decrypted server-side and never stored. + encrypted_api_key is decrypted server-side and never stored. """ info = get_provider_info(payload.provider_type) if info is None: @@ -220,9 +205,7 @@ async def test_provider( 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 - ) + 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.", @@ -254,10 +237,15 @@ async def test_provider( models_count = len(models), ) except Exception as exc: - logger.warning("Provider test failed for %s: %s", payload.provider_type, exc) + logger.error( + "providers.test_failed", + provider_type = payload.provider_type, + error = str(exc), + exc_info = True, + ) return ProviderTestResult( success = False, - message = f"Connection failed: {exc}", + message = f"Connection failed: {safe_curated_detail(exc)}", models_count = None, ) finally: @@ -269,13 +257,12 @@ async def test_provider( @router.post("/models", response_model = list[ProviderModelInfo]) async def list_provider_models( - payload: ProviderModelsRequest, - current_subject: str = Depends(get_current_subject), + payload: ProviderModelsRequest, current_subject: str = Depends(get_current_subject) ): """ List models available from an external provider. - The encrypted_api_key is decrypted server-side and never stored. + encrypted_api_key is decrypted server-side and never stored. """ info = get_provider_info(payload.provider_type) if info is None: @@ -289,9 +276,7 @@ async def list_provider_models( 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 - ) + 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.", @@ -318,27 +303,38 @@ async def list_provider_models( try: models = await client.list_models() - allow_prefixes = info.get("model_id_allow_prefixes") - if allow_prefixes is not None: - prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p)) - if prefix_tuple: - models = [m for m in models if m.get("id", "").startswith(prefix_tuple)] - allowlist = info.get("model_id_allowlist") - if allowlist is not None: - models = [m for m in models if allowlist.match(m.get("id", ""))] - deny_exact = info.get("model_id_deny_exact") - if deny_exact is not None: - deny_ids = {str(m) for m in deny_exact if str(m)} - if deny_ids: - models = [m for m in models if m.get("id", "") not in deny_ids] - denylist = info.get("model_id_denylist") - if denylist is not None: - models = [m for m in models if not denylist.search(m.get("id", ""))] - # Apply an optional cap after filtering so registry entries with a - # large remote catalog (e.g. HF Inference Providers) can stay - # picker-sized. No popularity sort happens server-side, so this is - # "first N matches" — pair with default_models for any must-have - # flagship ids. + # Registry model-id filters only apply to the native Gemini base. A + # custom OAI-compatible proxy returns prefixed IDs the native allowlist + # would strip, leaving the picker empty; match the host check here so the + # model list and chat dispatch agree on what counts as "native". + apply_registry_model_filters = True + if payload.provider_type == "gemini": + try: + from urllib.parse import urlparse as _urlparse + _host = (_urlparse(base_url).hostname or "").lower() + except Exception: + _host = "" + apply_registry_model_filters = _host == "generativelanguage.googleapis.com" + + if apply_registry_model_filters: + allow_prefixes = info.get("model_id_allow_prefixes") + if allow_prefixes is not None: + prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p)) + if prefix_tuple: + models = [m for m in models if m.get("id", "").startswith(prefix_tuple)] + allowlist = info.get("model_id_allowlist") + if allowlist is not None: + models = [m for m in models if allowlist.match(m.get("id", ""))] + deny_exact = info.get("model_id_deny_exact") + if deny_exact is not None: + deny_ids = {str(m) for m in deny_exact if str(m)} + if deny_ids: + models = [m for m in models if m.get("id", "") not in deny_ids] + denylist = info.get("model_id_denylist") + if denylist is not None: + models = [m for m in models if not denylist.search(m.get("id", ""))] + # Optional cap after filtering to keep large catalogs picker-sized. + # Unsorted, so "first N matches"; pair with default_models for flagships. limit = info.get("model_id_limit") if isinstance(limit, int) and limit > 0: models = models[:limit] @@ -352,10 +348,12 @@ async def list_provider_models( for m in models ] except Exception as exc: - logger.error("Failed to list models from %s: %s", payload.provider_type, exc) - raise HTTPException( - status_code = 502, - detail = f"Failed to list models from {payload.provider_type}: {exc}", + raise log_and_http_error( + exc, + 502, + f"Failed to list models from {payload.provider_type}.", + event = "providers.list_models_failed", + log = logger, ) finally: await client.close() diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py new file mode 100644 index 0000000000..8aba73caac --- /dev/null +++ b/studio/backend/routes/rag.py @@ -0,0 +1,456 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""HTTP API for the RAG engine: KB CRUD, uploads, SSE ingestion, search. + +Single-tenant: the subject gates access, not data. Without sqlite-vec the router +mounts but every endpoint returns 503. +""" + +from __future__ import annotations + +import hashlib +import hmac +import json +import logging +import os +import re +import secrets +import time +import uuid + +from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile +from fastapi.responses import FileResponse, StreamingResponse +from pydantic import BaseModel, Field + +from auth.authentication import get_current_subject +from core.rag import config, ingestion, retrieval, store +from storage import rag_db +from utils.paths import ensure_dir, rag_uploads_root + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +def _require_rag() -> None: + if not rag_db.RAG_AVAILABLE: + raise HTTPException( + status_code = 503, + detail = "RAG is unavailable: the sqlite-vec extension could not be loaded.", + ) + + +_SAFE = re.compile(r"[^A-Za-z0-9._-]+") + + +def _sanitize_filename(name: str) -> str: + base = os.path.basename(name or "").strip() or "document" + base = _SAFE.sub("_", base) + return base[:200] + + +def _save_upload(file: UploadFile) -> tuple[str, str]: + """Persist an upload; returns (stored_path, filename).""" + filename = _sanitize_filename(file.filename or "document") + ext = os.path.splitext(filename)[1].lower() + if ext not in config.UPLOAD_EXTS: + raise HTTPException( + status_code = 400, + detail = f"Unsupported file type '{ext}'. Allowed: {sorted(config.UPLOAD_EXTS)}", + ) + uploads = ensure_dir(rag_uploads_root()) + stored_path = str(uploads / f"{uuid.uuid4().hex}{ext}") + size = 0 + with open(stored_path, "wb") as out: + while True: + block = file.file.read(1 << 20) + if not block: + break + size += len(block) + out.write(block) + if size == 0: + os.remove(stored_path) + raise HTTPException(status_code = 400, detail = "Uploaded file is empty.") + return stored_path, filename + + +def _doc_view(row: dict) -> dict: + return { + "id": row["id"], + "filename": row["filename"], + "status": row["status"], + "error": row.get("error"), + "numChunks": row.get("num_chunks") or 0, + "kbId": row.get("kb_id"), + "threadId": row.get("thread_id"), + "createdAt": row.get("created_at"), + } + + +class CreateKbRequest(BaseModel): + name: str = Field(min_length = 1, max_length = 200) + description: str | None = None + + +class UpdateKbRequest(BaseModel): + name: str | None = Field(default = None, max_length = 200) + description: str | None = None + + +class SearchRequest(BaseModel): + query: str + kb_id: str | None = None + thread_id: str | None = None + top_k: int = Field(default = config.TOP_K_HYBRID, ge = 1, le = 50) + min_score: float = 0.0 + mode: str = "hybrid" # hybrid | lexical | dense + + +@router.get("/knowledge-bases") +def list_knowledge_bases(subject: str = Depends(get_current_subject)) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + kbs = store.list_kbs(conn) + out = [] + for kb in kbs: + docs = store.list_documents(conn, store.kb_scope(kb["id"])) + out.append( + { + "id": kb["id"], + "name": kb["name"], + "description": kb.get("description"), + "createdAt": kb.get("created_at"), + "documentCount": len(docs), + } + ) + return {"knowledgeBases": out} + finally: + conn.close() + + +@router.post("/knowledge-bases") +def create_knowledge_base( + payload: CreateKbRequest, subject: str = Depends(get_current_subject) +) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + kb_id = store.create_kb( + conn, + name = payload.name.strip(), + description = (payload.description or None), + embedding_model = config.EMBEDDING_MODEL, + ) + return {"id": kb_id, "name": payload.name.strip()} + finally: + conn.close() + + +@router.patch("/knowledge-bases/{kb_id}") +def update_knowledge_base( + kb_id: str, + payload: UpdateKbRequest, + subject: str = Depends(get_current_subject), +) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + if store.get_kb(conn, kb_id) is None: + raise HTTPException(status_code = 404, detail = "Knowledge base not found") + sets, params = [], [] + if payload.name is not None: + sets.append("name=?") + params.append(payload.name.strip()) + if payload.description is not None: + sets.append("description=?") + params.append(payload.description or None) + if sets: + params.append(kb_id) + conn.execute(f"UPDATE knowledge_bases SET {', '.join(sets)} WHERE id=?", params) + conn.commit() + return {"ok": True} + finally: + conn.close() + + +@router.delete("/knowledge-bases/{kb_id}") +def delete_knowledge_base(kb_id: str, subject: str = Depends(get_current_subject)) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + if store.get_kb(conn, kb_id) is None: + raise HTTPException(status_code = 404, detail = "Knowledge base not found") + store.delete_kb(conn, kb_id) + return {"ok": True} + finally: + conn.close() + + +@router.post("/knowledge-bases/{kb_id}/documents") +async def upload_kb_document( + kb_id: str, + file: UploadFile = File(...), + subject: str = Depends(get_current_subject), +) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + if store.get_kb(conn, kb_id) is None: + raise HTTPException(status_code = 404, detail = "Knowledge base not found") + finally: + conn.close() + stored_path, filename = _save_upload(file) + document_id, job_id = ingestion.start_ingestion( + store.kb_scope(kb_id), kb_id, None, filename, stored_path + ) + return {"documentId": document_id, "jobId": job_id, "filename": filename} + + +@router.get("/knowledge-bases/{kb_id}/documents") +def list_kb_documents(kb_id: str, subject: str = Depends(get_current_subject)) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + docs = store.list_documents(conn, store.kb_scope(kb_id)) + return {"documents": [_doc_view(d) for d in docs]} + finally: + conn.close() + + +@router.post("/threads/{thread_id}/documents") +async def upload_thread_document( + thread_id: str, + file: UploadFile = File(...), + subject: str = Depends(get_current_subject), +) -> dict: + _require_rag() + stored_path, filename = _save_upload(file) + document_id, job_id = ingestion.start_ingestion( + store.thread_scope(thread_id), None, thread_id, filename, stored_path + ) + return {"documentId": document_id, "jobId": job_id, "filename": filename} + + +@router.get("/threads/{thread_id}/documents") +def list_thread_documents(thread_id: str, subject: str = Depends(get_current_subject)) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + docs = store.list_documents(conn, store.thread_scope(thread_id)) + return {"documents": [_doc_view(d) for d in docs]} + finally: + conn.close() + + +@router.delete("/documents/{document_id}") +def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + if store.get_document(conn, document_id) is None: + raise HTTPException(status_code = 404, detail = "Document not found") + store.delete_document(conn, document_id) + return {"ok": True} + finally: + conn.close() + + +@router.get("/jobs/{job_id}") +def job_status(job_id: str, subject: str = Depends(get_current_subject)) -> dict: + _require_rag() + row = ingestion.get_job_status(job_id) + if row is None: + raise HTTPException(status_code = 404, detail = "Job not found") + return { + "id": row["id"], + "documentId": row["document_id"], + "status": row["status"], + "stage": row.get("stage"), + "progress": row.get("progress") or 0.0, + "error": row.get("error"), + } + + +@router.get("/jobs/{job_id}/events") +def job_events(job_id: str, subject: str = Depends(get_current_subject)) -> StreamingResponse: + _require_rag() + + def gen(): + try: + for event in ingestion.job_events(job_id): + yield f"data: {json.dumps(event)}\n\n" + except Exception as exc: # noqa: BLE001 + yield f"data: {json.dumps({'type': 'error', 'error': str(exc)})}\n\n" + yield "data: [DONE]\n\n" + + return StreamingResponse( + gen(), + media_type = "text/event-stream", + headers = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + ) + + +@router.post("/search") +def search(payload: SearchRequest, subject: str = Depends(get_current_subject)) -> dict: + _require_rag() + if payload.kb_id: + scope = store.kb_scope(payload.kb_id) + elif payload.thread_id: + scope = store.thread_scope(payload.thread_id) + else: + raise HTTPException(status_code = 400, detail = "Provide kb_id or thread_id") + + conn = rag_db.get_connection() + try: + if payload.mode == "lexical": + hits = retrieval.retrieve_lexical(conn, scope, payload.query, payload.top_k) + elif payload.mode == "dense": + hits = retrieval.retrieve_dense(conn, scope, payload.query, payload.top_k) + else: + hits = retrieval.retrieve_hybrid(conn, scope, payload.query, k = payload.top_k) + hits = retrieval.filter_min_score(hits, payload.min_score) + rows = store.chunks_by_id(conn, [h.chunk_id for h in hits]) + results = [] + for h in hits: + r = rows.get(h.chunk_id) + if r is None: + continue + results.append( + { + "chunkId": h.chunk_id, + "documentId": r["document_id"], + "filename": r["filename"], + "page": r["page_number"], + "score": h.score, + "text": r["text"], + } + ) + return {"results": results} + finally: + conn.close() + + +# Per-process secret so pdf.js range requests fetch the file without a bearer +# header; tokens only work on this server instance. +_PREVIEW_SECRET = secrets.token_bytes(32) +_PREVIEW_TTL = 600 # seconds + +_CONTENT_TYPES = { + ".pdf": "application/pdf", + ".txt": "text/plain; charset=utf-8", + ".md": "text/markdown; charset=utf-8", + ".markdown": "text/markdown; charset=utf-8", + ".html": "text/html; charset=utf-8", + ".htm": "text/html; charset=utf-8", + ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", +} + + +def _sign_document(document_id: str) -> str: + exp = int(time.time()) + _PREVIEW_TTL + payload = f"{document_id}.{exp}" + sig = hmac.new(_PREVIEW_SECRET, payload.encode(), hashlib.sha256).hexdigest() + return f"{payload}.{sig}" + + +def _verify_document_token(token: str) -> str | None: + try: + document_id, exp_s, sig = token.rsplit(".", 2) + except ValueError: + return None + expected = hmac.new( + _PREVIEW_SECRET, f"{document_id}.{exp_s}".encode(), hashlib.sha256 + ).hexdigest() + if not hmac.compare_digest(sig, expected): + return None + try: + if int(exp_s) < int(time.time()): + return None + except ValueError: + return None + return document_id + + +@router.get("/documents/{document_id}/preview-target") +def preview_target( + document_id: str, + chunk_id: str | None = Query(default = None), + subject: str = Depends(get_current_subject), +) -> dict: + """Resolve a citation to filename, page, and highlight regions.""" + _require_rag() + conn = rag_db.get_connection() + try: + doc = store.get_document(conn, document_id) + if doc is None: + raise HTTPException(status_code = 404, detail = "Document not found") + ext = os.path.splitext(doc["filename"])[1].lower() + out = { + "documentId": document_id, + "filename": doc["filename"], + "mediaKind": "pdf" if ext == ".pdf" else "text", + "targetPage": None, + "pdfRegions": [], + "text": None, + } + if chunk_id: + row = conn.execute( + "SELECT text, page_number, pdf_regions_json FROM chunks WHERE id=?", + (chunk_id,), + ).fetchone() + if row is not None: + out["text"] = row["text"] + out["targetPage"] = row["page_number"] + if row["pdf_regions_json"]: + try: + out["pdfRegions"] = json.loads(row["pdf_regions_json"]) + except Exception: + out["pdfRegions"] = [] + return out + finally: + conn.close() + + +@router.get("/documents/{document_id}/file-url") +def document_file_url(document_id: str, subject: str = Depends(get_current_subject)) -> dict: + """Mint a short-lived signed URL for the source file.""" + _require_rag() + conn = rag_db.get_connection() + try: + doc = store.get_document(conn, document_id) + if doc is None or not doc.get("stored_path"): + raise HTTPException(status_code = 404, detail = "Document file not available") + finally: + conn.close() + token = _sign_document(document_id) + return {"url": f"/api/rag/documents/{document_id}/file-signed?token={token}"} + + +@router.get("/documents/{document_id}/file-signed", response_model = None) +def document_file_signed(document_id: str, token: str = Query(...)) -> FileResponse: + """Serve the source file gated by the HMAC token (no bearer) so pdf.js range + requests work.""" + _require_rag() + signed_id = _verify_document_token(token) + if signed_id != document_id: + raise HTTPException(status_code = 401, detail = "Invalid or expired token") + conn = rag_db.get_connection() + try: + doc = store.get_document(conn, document_id) + finally: + conn.close() + stored_path = (doc or {}).get("stored_path") + if not doc or not stored_path or not os.path.isfile(stored_path): + raise HTTPException(status_code = 404, detail = "Document file not found") + # Confine to the uploads root (defense in depth). + uploads = os.path.realpath(str(rag_uploads_root())) + if not os.path.realpath(stored_path).startswith(uploads): + raise HTTPException(status_code = 403, detail = "Forbidden") + ext = os.path.splitext(doc["filename"])[1].lower() + return FileResponse( + stored_path, + media_type = _CONTENT_TYPES.get(ext, "application/octet-stream"), + filename = doc["filename"], + ) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py new file mode 100644 index 0000000000..516502c1e5 --- /dev/null +++ b/studio/backend/routes/settings.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from fastapi import APIRouter, Depends +from pydantic import BaseModel, Field + +from auth.authentication import get_current_subject +from loggers import get_logger +from utils.utils import safe_error_detail, log_and_http_error +from utils.upload_limits import ( + MAX_UPLOAD_LIMIT_MB, + MIN_UPLOAD_LIMIT_MB, + default_upload_limit_mb, + get_upload_limit_mb, + set_upload_limit_mb, + upload_limit_bytes, + upload_limit_label, +) +from utils.helper_precache_settings import ( + DEFAULT_HELPER_PRECACHE_ENABLED, + get_helper_precache_enabled, + helper_model_disabled_by_env, + set_helper_precache_enabled, +) + +router = APIRouter() + +logger = get_logger(__name__) + + +class UploadLimitPayload(BaseModel): + max_upload_size_mb: int = Field(..., ge = MIN_UPLOAD_LIMIT_MB, le = MAX_UPLOAD_LIMIT_MB) + + +class UploadLimitResponse(BaseModel): + max_upload_size_mb: int + max_upload_size_bytes: int + max_upload_size_label: str + default_upload_size_mb: int + min_upload_size_mb: int = MIN_UPLOAD_LIMIT_MB + max_allowed_upload_size_mb: int = MAX_UPLOAD_LIMIT_MB + + +class HelperPrecachePayload(BaseModel): + enabled: bool + + +class HelperPrecacheResponse(BaseModel): + enabled: bool + default_enabled: bool = DEFAULT_HELPER_PRECACHE_ENABLED + disabled_by_env: bool + + +def _upload_limit_response(limit_mb: int) -> UploadLimitResponse: + return UploadLimitResponse( + max_upload_size_mb = limit_mb, + max_upload_size_bytes = upload_limit_bytes(limit_mb), + max_upload_size_label = upload_limit_label(limit_mb), + default_upload_size_mb = default_upload_limit_mb(), + ) + + +def _helper_precache_response(enabled: bool | None = None) -> HelperPrecacheResponse: + return HelperPrecacheResponse( + enabled = get_helper_precache_enabled() if enabled is None else enabled, + disabled_by_env = helper_model_disabled_by_env(), + ) + + +@router.get("/upload-limit", response_model = UploadLimitResponse) +def get_upload_limit(current_subject: str = Depends(get_current_subject)) -> UploadLimitResponse: + return _upload_limit_response(get_upload_limit_mb()) + + +@router.put("/upload-limit", response_model = UploadLimitResponse) +def update_upload_limit( + payload: UploadLimitPayload, current_subject: str = Depends(get_current_subject) +) -> UploadLimitResponse: + try: + limit_mb = set_upload_limit_mb(payload.max_upload_size_mb) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid upload limit."), + event = "settings.update_upload_limit_failed", + log = logger, + ) from exc + return _upload_limit_response(limit_mb) + + +@router.get("/helper-precache", response_model = HelperPrecacheResponse) +def get_helper_precache( + current_subject: str = Depends(get_current_subject), +) -> HelperPrecacheResponse: + return _helper_precache_response() + + +@router.put("/helper-precache", response_model = HelperPrecacheResponse) +def update_helper_precache( + payload: HelperPrecachePayload, current_subject: str = Depends(get_current_subject) +) -> HelperPrecacheResponse: + try: + enabled = set_helper_precache_enabled(payload.enabled) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid Helper LLM pre-cache setting."), + event = "settings.update_helper_precache_failed", + log = logger, + ) from exc + return _helper_precache_response(enabled) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 6e2413b3e9..d7687ffdee 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -16,13 +16,11 @@ import asyncio from datetime import datetime import uuid as _uuid -# Add backend directory to path -# The backend code should be in the same directory structure +# Add backend directory to path. backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) -# Import backend functions try: from core.training import get_training_backend from core.training.resume import ( @@ -34,7 +32,7 @@ try: from utils.models.model_config import load_model_defaults from utils.paths import resolve_dataset_path except ImportError: - # Fallback: try to import from parent directory + # Fallback: parent directory. parent_backend = backend_path.parent / "backend" if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) @@ -51,6 +49,8 @@ except ImportError: # Auth from auth.authentication import get_current_subject +from utils.utils import log_and_http_error + from models import ( TrainingStartRequest, TrainingJobResponse, @@ -69,9 +69,7 @@ router = APIRouter() logger = get_logger(__name__) -def _validate_local_dataset_paths( - paths: list[str], label: str = "Local dataset" -) -> list[str]: +def _validate_local_dataset_paths(paths: list[str], label: str = "Local dataset") -> list[str]: """Resolve and validate a list of local dataset paths. Returns validated absolute paths.""" validated = [] missing = [] @@ -93,50 +91,41 @@ def _validate_local_dataset_paths( @router.get("/hardware") -async def get_hardware_utilization( - current_subject: str = Depends(get_current_subject), -): +async def get_hardware_utilization(current_subject: str = Depends(get_current_subject)): """ - Get a live snapshot of GPU hardware utilization. + Live snapshot of GPU hardware utilization for the active backend. - Designed to be polled by the frontend during training. - Returns live GPU memory usage information for the active backend. + Polled by the frontend during training. """ from utils.hardware import get_gpu_utilization - return get_gpu_utilization() @router.get("/hardware/visible") -async def get_visible_hardware_utilization( - current_subject: str = Depends(get_current_subject), -): +async def get_visible_hardware_utilization(current_subject: str = Depends(get_current_subject)): from utils.hardware import get_visible_gpu_utilization - return get_visible_gpu_utilization() @router.post("/start") async def start_training( - request: TrainingStartRequest, - current_subject: str = Depends(get_current_subject), + request: TrainingStartRequest, current_subject: str = Depends(get_current_subject) ): """ Start a training job. - This endpoint initiates training in the background and returns immediately. - Use the /status endpoint to check training progress. + Initiates training in the background and returns immediately. Use /status + to check progress. """ try: logger.info(f"Starting training job with model: {request.model_name}") - # NOTE: No in-process ensure_transformers_version() call here. - # The subprocess (worker.py) activates the correct version in a - # fresh Python interpreter before importing any ML libraries. + # No in-process ensure_transformers_version(): the subprocess + # (worker.py) activates the correct version before importing ML libs. backend = get_training_backend() - # Check if training is already active (before mutating any state) + # Check before mutating state. if backend.is_training_active(): existing_job_id: Optional[str] = getattr(backend, "current_job_id", "") return TrainingJobResponse( @@ -149,13 +138,11 @@ async def start_training( error = "Training already active", ) - # Generate job ID — passed into start_training() which sets it on the - # backend only after confirming the old pump thread is dead. - job_id = ( - f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}" - ) + # Job ID; start_training() sets it on the backend only after the old + # pump thread is dead. + job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}" - # Validate dataset paths if provided + # Validate dataset paths if provided. if request.local_datasets: request.local_datasets = _validate_local_dataset_paths( request.local_datasets, "Local dataset" @@ -167,11 +154,11 @@ async def start_training( resume_output_dir: Optional[str] = None if request.resume_from_checkpoint: try: - resume_output_dir = normalize_resume_output_dir( - request.resume_from_checkpoint - ) + resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint) except ValueError as e: - raise HTTPException(status_code = 400, detail = str(e)) + # Deliberate user-facing validation message. + validation_message = str(e) + raise HTTPException(status_code = 400, detail = validation_message) resume_run = get_resumable_run_by_output_dir(resume_output_dir) if not resume_run or not can_resume_run(resume_run): @@ -187,13 +174,14 @@ async def start_training( ) request.resume_from_checkpoint = resume_checkpoint - # Convert request to kwargs for backend + # Convert request to backend kwargs. training_kwargs = { "model_name": request.model_name, "training_type": request.training_type, "hf_token": request.hf_token or "", "load_in_4bit": request.load_in_4bit, "max_seq_length": request.max_seq_length, + "vision_image_size": request.vision_image_size, "hf_dataset": request.hf_dataset or "", "local_datasets": request.local_datasets, "local_eval_datasets": request.local_eval_datasets, @@ -224,9 +212,7 @@ async def start_training( "lora_r": request.lora_r, "lora_alpha": request.lora_alpha, "lora_dropout": request.lora_dropout, - "target_modules": request.target_modules - if request.target_modules - else None, + "target_modules": request.target_modules if request.target_modules else None, "gradient_checkpointing": request.gradient_checkpointing.strip() if request.gradient_checkpointing and request.gradient_checkpointing.strip() else "unsloth", @@ -251,25 +237,19 @@ async def start_training( "gpu_ids": request.gpu_ids, } - # Training page has no trust_remote_code toggle — the value comes from - # YAML model defaults applied when the user selects a model. As a safety - # net, consult the YAML directly so models that need it always get it. + # Training page has no trust_remote_code toggle; as a safety net consult + # YAML model defaults directly so models that need it always get it. if not training_kwargs["trust_remote_code"]: model_defaults = load_model_defaults(request.model_name) - yaml_trust = model_defaults.get("training", {}).get( - "trust_remote_code", False - ) + yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False) if yaml_trust: - logger.info( - f"YAML config sets trust_remote_code=True for {request.model_name}" - ) + logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}") training_kwargs["trust_remote_code"] = True # Free GPU memory: shut down any running inference/export subprocesses - # before training starts (they'd compete for VRAM otherwise) + # before training (they'd compete for VRAM otherwise). try: from core.inference import get_inference_backend - inf_backend = get_inference_backend() if inf_backend.active_model_name: logger.info( @@ -284,12 +264,9 @@ async def start_training( try: from core.export import get_export_backend - exp_backend = get_export_backend() if exp_backend.current_checkpoint: - logger.info( - "Shutting down export subprocess to free GPU memory for training" - ) + logger.info("Shutting down export subprocess to free GPU memory for training") exp_backend._shutdown_subprocess() exp_backend.current_checkpoint = None exp_backend.is_vision = False @@ -297,7 +274,7 @@ async def start_training( except Exception as e: logger.warning("Could not shut down export subprocess: %s", e) - # start_training now spawns a subprocess (non-blocking) + # start_training spawns a subprocess (non-blocking). success = backend.start_training(job_id = job_id, **training_kwargs) if not success: @@ -318,12 +295,16 @@ async def start_training( except ValueError as e: logger.warning("Rejected training GPU selection: %s", e) - raise HTTPException(status_code = 400, detail = str(e)) + # Deliberate user-facing GPU-selection validation message. + validation_message = str(e) + raise HTTPException(status_code = 400, detail = validation_message) except Exception as e: - logger.error(f"Error starting training: {e}", exc_info = True) - raise HTTPException( - status_code = 500, - detail = f"Failed to start training: {str(e)}", + raise log_and_http_error( + e, + 500, + "Failed to start training", + event = "training.start_failed", + log = logger, ) @@ -348,7 +329,6 @@ async def stop_training( status = "idle", message = "No training job is currently running" ) - # Call backend stop method backend.stop_training(save = body.save) return TrainingStopResponse( @@ -357,34 +337,29 @@ async def stop_training( ) except Exception as e: - logger.error(f"Error stopping training: {e}", exc_info = True) - raise HTTPException( - status_code = 500, detail = f"Failed to stop training: {str(e)}" + raise log_and_http_error( + e, + 500, + "Failed to stop training", + event = "training.stop_failed", + log = logger, ) @router.post("/reset") -async def reset_training( - current_subject: str = Depends(get_current_subject), -): - """ - Reset training state so the user can return to configuration. - """ +async def reset_training(current_subject: str = Depends(get_current_subject)): + """Reset training state so the user can return to configuration.""" try: backend = get_training_backend() is_active = backend.is_training_active() if is_active: if backend._cancel_requested: - # Cancel (save=False) was requested — force-terminate so we can reset immediately - logger.info( - "Force-terminating subprocess for immediate reset (cancel path)" - ) + # Cancel (save=False) requested — force-terminate to reset immediately. + logger.info("Force-terminating subprocess for immediate reset (cancel path)") backend.force_terminate() else: - logger.warning( - "Rejected reset while training active: is_active=%s", is_active - ) + logger.warning("Rejected reset while training active: is_active=%s", is_active) raise HTTPException( status_code = 409, detail = "Training is still running. Stop training and wait for it to finish before resetting.", @@ -411,17 +386,17 @@ async def reset_training( except HTTPException: raise except Exception as e: - logger.error(f"Error resetting training: {e}", exc_info = True) - raise HTTPException( - status_code = 500, - detail = f"Failed to reset training: {str(e)}", + raise log_and_http_error( + e, + 500, + "Failed to reset training", + event = "training.reset_failed", + log = logger, ) @router.get("/status") -async def get_training_status( - current_subject: str = Depends(get_current_subject), -): +async def get_training_status(current_subject: str = Depends(get_current_subject)): """ Get the current training status. """ @@ -429,10 +404,8 @@ async def get_training_status( backend = get_training_backend() job_id: str = getattr(backend, "current_job_id", "") or "" - # Check if training is active is_active = backend.is_training_active() - # Get progress info from trainer try: progress = backend.trainer.get_training_progress() except Exception: @@ -443,7 +416,6 @@ async def get_training_status( ) or "Ready to train" error_message = getattr(progress, "error", None) if progress else None - # Check if training was stopped by user trainer_stopped = getattr(backend, "_should_stop", False) # Derive high-level phase @@ -453,9 +425,7 @@ async def get_training_status( msg_lower = status_message.lower() if "loading" in msg_lower or "importing" in msg_lower: phase = "loading_model" - elif any( - k in msg_lower for k in ["preparing", "initializing", "configuring"] - ): + elif any(k in msg_lower for k in ["preparing", "initializing", "configuring"]): phase = "configuring" else: phase = "training" @@ -479,7 +449,7 @@ async def get_training_status( if output_dir: details["output_dir"] = output_dir - # Build metric history for chart recovery after SSE reconnection + # Metric history for chart recovery after SSE reconnection. metric_history = None if backend.step_history: metric_history = { @@ -504,30 +474,29 @@ async def get_training_status( ) except Exception as e: - logger.error(f"Error getting training status: {e}", exc_info = True) - raise HTTPException( - status_code = 500, detail = f"Failed to get training status: {str(e)}" + raise log_and_http_error( + e, + 500, + "Failed to get training status", + event = "training.status_failed", + log = logger, ) @router.get("/metrics", response_model = TrainingMetricsResponse) -async def get_training_metrics( - current_subject: str = Depends(get_current_subject), -): +async def get_training_metrics(current_subject: str = Depends(get_current_subject)): """ Get training metrics (loss, learning rate, steps). """ try: backend = get_training_backend() - # Get metrics from backend loss_history = backend.loss_history lr_history = backend.lr_history step_history = backend.step_history grad_norm_history = getattr(backend, "grad_norm_history", []) grad_norm_step_history = getattr(backend, "grad_norm_step_history", []) - # Get current values current_loss = loss_history[-1] if loss_history else None current_lr = lr_history[-1] if lr_history else None current_step = step_history[-1] if step_history else None @@ -544,28 +513,29 @@ async def get_training_metrics( ) except Exception as e: - logger.error(f"Error getting training metrics: {e}", exc_info = True) - raise HTTPException( - status_code = 500, detail = f"Failed to get training metrics: {str(e)}" + raise log_and_http_error( + e, + 500, + "Failed to get training metrics", + event = "training.metrics_failed", + log = logger, ) @router.get("/progress") async def stream_training_progress( - request: Request, - current_subject: str = Depends(get_current_subject), + request: Request, current_subject: str = Depends(get_current_subject) ): """ - Stream training progress updates using Server-Sent Events (SSE). + Stream training progress via Server-Sent Events (SSE). - This endpoint provides real-time updates on training progress. - Supports reconnection via the SSE spec: - - Sends `id:` with each event so the browser tracks position. - - Sends `retry:` to control reconnection interval. - - Sends named `event:` types (progress, heartbeat, complete, error). - - Reads `Last-Event-ID` header on reconnect to replay missed steps. + Real-time progress with reconnection support per the SSE spec: + - `id:` per event so the browser tracks position. + - `retry:` to control reconnection interval. + - Named `event:` types (progress, heartbeat, complete, error). + - Reads `Last-Event-ID` on reconnect to replay missed steps. """ - # Read Last-Event-ID header for reconnection resume + # Read Last-Event-ID header for reconnection resume. last_event_id = request.headers.get("last-event-id") resume_from_step: Optional[int] = None if last_event_id is not None: @@ -594,14 +564,10 @@ async def stream_training_progress( if step < 0 or total == 0: progress_percent = 0.0 else: - progress_percent = ( - float(step) / float(total) * 100.0 if total > 0 else 0.0 - ) + progress_percent = float(step) / float(total) * 100.0 if total > 0 else 0.0 - # Get actual values from progress object if available - elapsed_seconds = ( - getattr(progress, "elapsed_seconds", None) if progress else None - ) + # Pull values from the progress object if available. + elapsed_seconds = getattr(progress, "elapsed_seconds", None) if progress else None eta_seconds = getattr(progress, "eta_seconds", None) if progress else None grad_norm = grad_norm_override if grad_norm is None and progress: @@ -642,7 +608,7 @@ async def stream_training_progress( return "\n".join(lines) # ── Retry directive ────────────────────────────────────── - # Tell the browser to reconnect after 3 seconds if the connection drops + # Reconnect after 3 seconds if the connection drops. yield "retry: 3000\n\n" # ── Replay missed steps on reconnect ───────────────────── @@ -657,25 +623,15 @@ async def stream_training_progress( } for i, step_val in enumerate(backend.step_history): if step_val > resume_from_step: - loss_val = ( - backend.loss_history[i] - if i < len(backend.loss_history) - else None - ) - lr_val = ( - backend.lr_history[i] if i < len(backend.lr_history) else None - ) + loss_val = backend.loss_history[i] if i < len(backend.loss_history) else None + lr_val = backend.lr_history[i] if i < len(backend.lr_history) else None tp_replay = getattr( getattr(backend, "trainer", None), "training_progress", None ) total_replay = ( - getattr(tp_replay, "total_steps", step_val) - if tp_replay - else step_val - ) - epoch_replay = ( - getattr(tp_replay, "epoch", None) if tp_replay else None + getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val ) + epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None payload = build_progress( step_val, loss_val, @@ -685,9 +641,7 @@ async def stream_training_progress( progress = tp_replay, grad_norm_override = grad_norm_by_step.get(step_val), ) - yield format_sse( - payload.model_dump_json(), event = "progress", event_id = step_val - ) + yield format_sse(payload.model_dump_json(), event = "progress", event_id = step_val) replayed += 1 if replayed: logger.info(f"SSE reconnect: replayed {replayed} missed steps") @@ -707,21 +661,22 @@ async def stream_training_progress( epoch = initial_epoch, progress = tp, ) - yield format_sse( - initial_progress.model_dump_json(), event = "progress", event_id = 0 - ) + yield format_sse(initial_progress.model_dump_json(), event = "progress", event_id = 0) # If not active, send final state and exit if not is_active: - if backend.step_history: - final_step = backend.step_history[-1] - final_loss = ( - backend.loss_history[-1] if backend.loss_history else None - ) + _live = (getattr(tp, "step", 0) or 0) if tp else 0 + if backend.step_history or _live > 0: + final_step = backend.step_history[-1] if backend.step_history else 0 + final_loss = backend.loss_history[-1] if backend.loss_history else None final_lr = backend.lr_history[-1] if backend.lr_history else None - final_total_steps = ( - getattr(tp, "total_steps", final_step) if tp else final_step - ) + # Histories skip non-finite steps; report the live step with + # loss=None instead of the last finite pair. + if _live > final_step: + final_step = _live + final_loss = getattr(tp, "loss", None) + final_lr = getattr(tp, "learning_rate", final_lr) + final_total_steps = getattr(tp, "total_steps", final_step) if tp else final_step final_epoch = getattr(tp, "epoch", None) if tp else None payload = build_progress( final_step, @@ -736,9 +691,7 @@ async def stream_training_progress( ) else: yield format_sse( - build_progress( - -1, None, None, 0, progress = tp - ).model_dump_json(), + build_progress(-1, None, None, 0, progress = tp).model_dump_json(), event = "complete", event_id = 0, ) @@ -747,31 +700,28 @@ async def stream_training_progress( # ── Live polling loop ──────────────────────────────────── last_step = resume_from_step if resume_from_step is not None else -1 no_update_count = 0 - max_no_updates = ( - 1800 # Timeout after 30 minutes (large models need time for compilation) - ) + max_no_updates = 1800 # Timeout after 30 min (large models need compile time) while backend.is_training_active(): try: - if backend.step_history: - current_step = backend.step_history[-1] - current_loss = ( - backend.loss_history[-1] if backend.loss_history else None - ) + tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None) + live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0 + if backend.step_history or live_step > 0: + current_step = backend.step_history[-1] if backend.step_history else 0 + current_loss = backend.loss_history[-1] if backend.loss_history else None current_lr = backend.lr_history[-1] if backend.lr_history else None - tp_inner = getattr( - getattr(backend, "trainer", None), "training_progress", None - ) + # Histories skip non-finite steps; follow the live progress + # step and report its loss (None until it recovers). + if live_step > current_step: + current_step = live_step + current_loss = getattr(tp_inner, "loss", None) + current_lr = getattr(tp_inner, "learning_rate", current_lr) current_total_steps = ( - getattr(tp_inner, "total_steps", current_step) - if tp_inner - else current_step - ) - current_epoch = ( - getattr(tp_inner, "epoch", None) if tp_inner else None + getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step ) + current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None - # Only send if step changed + # Only send if the step changed. if current_step != last_step: progress_payload = build_progress( current_step, @@ -790,7 +740,7 @@ async def stream_training_progress( no_update_count = 0 else: no_update_count += 1 - # Send heartbeat every 10 seconds + # Heartbeat every 10 seconds. if no_update_count % 10 == 0: heartbeat_payload = build_progress( current_step, @@ -806,19 +756,17 @@ async def stream_training_progress( event_id = current_step, ) else: - # No steps yet, but training is active (model loading, etc.) + # No steps yet, but training is active (model loading, etc.). no_update_count += 1 if no_update_count % 5 == 0: - # Pull total_steps and status from trainer so - # the frontend can show "Tokenizing…" etc. + # Pull total_steps + status so the frontend can show + # "Tokenizing…" etc. tp_prep = getattr( getattr(backend, "trainer", None), "training_progress", None, ) - prep_total = ( - getattr(tp_prep, "total_steps", 0) if tp_prep else 0 - ) + prep_total = getattr(tp_prep, "total_steps", 0) if tp_prep else 0 preparing_payload = build_progress( 0, None, @@ -838,9 +786,7 @@ async def stream_training_progress( tp_timeout = getattr( getattr(backend, "trainer", None), "training_progress", None ) - timeout_payload = build_progress( - last_step, None, None, 0, progress = tp_timeout - ) + timeout_payload = build_progress(last_step, None, None, 0, progress = tp_timeout) yield format_sse( timeout_payload.model_dump_json(), event = "error", @@ -852,9 +798,7 @@ async def stream_training_progress( except Exception as e: logger.error(f"Error in progress stream: {e}", exc_info = True) - tp_error = getattr( - getattr(backend, "trainer", None), "training_progress", None - ) + tp_error = getattr(getattr(backend, "trainer", None), "training_progress", None) error_payload = build_progress(0, None, None, 0, progress = tp_error) yield format_sse( error_payload.model_dump_json(), @@ -868,9 +812,14 @@ async def stream_training_progress( final_loss = backend.loss_history[-1] if backend.loss_history else None final_lr = backend.lr_history[-1] if backend.lr_history else None final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None) - final_total_steps = ( - getattr(final_tp, "total_steps", final_step) if final_tp else final_step - ) + # If the run ended on a non-finite stretch, report the live step with + # loss=None instead of rolling back to the last finite pair. + _final_live_step = (getattr(final_tp, "step", 0) or 0) if final_tp else 0 + if _final_live_step > (final_step if final_step is not None else -1): + final_step = _final_live_step + final_loss = getattr(final_tp, "loss", None) + final_lr = getattr(final_tp, "learning_rate", final_lr) + final_total_steps = getattr(final_tp, "total_steps", final_step) if final_tp else final_step final_epoch = getattr(final_tp, "epoch", None) if final_tp else None final_payload = build_progress( final_step, diff --git a/studio/backend/routes/training_history.py b/studio/backend/routes/training_history.py index 771d9f1e35..1560c72767 100644 --- a/studio/backend/routes/training_history.py +++ b/studio/backend/routes/training_history.py @@ -42,19 +42,13 @@ async def list_training_runs( """List training runs, newest first.""" result = list_runs(limit = limit, offset = offset) return TrainingRunListResponse( - runs = [ - TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) - for r in result["runs"] - ], + runs = [TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) for r in result["runs"]], total = result["total"], ) @router.get("/runs/{run_id}", response_model = TrainingRunDetailResponse) -async def get_training_run_detail( - run_id: str, - current_subject: str = Depends(get_current_subject), -): +async def get_training_run_detail(run_id: str, current_subject: str = Depends(get_current_subject)): """Get a single training run with full config and metrics.""" run = get_run(run_id) if run is None: @@ -109,18 +103,13 @@ async def update_training_run( @router.delete("/runs/{run_id}", response_model = TrainingRunDeleteResponse) -async def delete_training_run( - run_id: str, - current_subject: str = Depends(get_current_subject), -): +async def delete_training_run(run_id: str, current_subject: str = Depends(get_current_subject)): """Delete a training run and its metrics (CASCADE).""" run = get_run(run_id) if run is None: raise HTTPException(status_code = 404, detail = f"Run {run_id} not found") if run["status"] == "running": - raise HTTPException( - status_code = 409, detail = "Cannot delete a running training run" - ) + raise HTTPException(status_code = 409, detail = "Cannot delete a running training run") logger.info("Deleting training run %s", run_id) delete_run(run_id) return TrainingRunDeleteResponse( diff --git a/studio/backend/run.py b/studio/backend/run.py index e96e609659..a883154c4a 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1,26 +1,35 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Run script for Unsloth UI Backend. -Works independently and can be moved to any directory. +"""Run script for Unsloth UI Backend. + +Self-contained; can be moved to any directory. """ import os import sys +import time from pathlib import Path from typing import Optional -# Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked) +# Suppress C-level dependency warnings globally (e.g. SwigPyPacked). os.environ["PYTHONWARNINGS"] = "ignore" -# Add the backend directory to Python path early so local modules are importable +# Add the backend dir to sys.path early so local modules import. backend_dir = Path(__file__).parent if str(backend_dir) not in sys.path: sys.path.insert(0, str(backend_dir)) -# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before -# any library imports that trigger attrs -> rich -> structlog -> platform crash. +from utils.cpu_threads import configure_cpu_threads + +try: + configure_cpu_threads() +except ValueError as exc: + configured = os.environ.get("UNSLOTH_CPU_THREADS") + raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}") from None + +# Anaconda/conda-forge Python: seed platform._sys_version_cache before imports +# that trigger attrs -> rich -> structlog -> platform crash. # See: https://github.com/python/cpython/issues/102396 import _platform_compat # noqa: F401 @@ -31,18 +40,17 @@ logger = get_logger(__name__) def _resolve_external_ip() -> str: - """ - Resolve the machine's external IP address. + """Resolve the machine's external IP address. - Tries (in order): - 1. GCE metadata server (instant, works on Google Cloud VMs) - 2. ifconfig.me (works anywhere with internet) + Tries, in order: + 1. GCE metadata server (instant on Google Cloud VMs) + 2. ifconfig.me (anywhere with internet) 3. LAN IP via UDP socket trick (fallback) """ import urllib.request import socket - # 1. Try GCE metadata server (responds in <10ms on GCE, times out fast elsewhere) + # 1. GCE metadata server (<10ms on GCE, times out fast elsewhere). try: req = urllib.request.Request( "http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/access-configs/0/external-ip", @@ -55,7 +63,7 @@ def _resolve_external_ip() -> str: except Exception: pass - # 2. Try public IP service + # 2. Public IP service. try: with urllib.request.urlopen("https://ifconfig.me", timeout = 3) as resp: ip = resp.read().decode().strip() @@ -77,15 +85,13 @@ def _resolve_external_ip() -> str: def _install_uvicorn_startup_log_rewrite(bind_host: str, display_host: str) -> None: """Rewrite Uvicorn's startup log line: swap wildcard bind for the - externally-reachable address, replace the CTRL+C suffix with our Mac-aware - stop hint, and rename the prefix to "Unsloth Studio running on".""" + externally-reachable address, use our Mac-aware stop hint, and rename the + prefix to "Unsloth Studio running on".""" import logging import re rewrite_host = ( - bind_host in ("0.0.0.0", "::") - and bool(display_host) - and display_host != bind_host + bind_host in ("0.0.0.0", "::") and bool(display_host) and display_host != bind_host ) new_suffix = "(To stop: press Ctrl+C -- on macOS, Control+C not Command+C)" old_suffix_re = re.compile(r"\(Press CTRL\+C to quit\)") @@ -126,10 +132,13 @@ def _install_uvicorn_startup_log_rewrite(bind_host: str, display_host: str) -> N logging.getLogger(name).addFilter(f) -def _local_port_open(host: str, port: int, timeout: float = 1.0) -> bool: - """Return True iff a TCP connection to (host, port) succeeds within timeout.""" +def _local_port_open( + host: str, + port: int, + timeout: float = 1.0, +) -> bool: + """True iff a TCP connection to (host, port) succeeds within timeout.""" import socket - try: with socket.create_connection((host, port), timeout = timeout): return True @@ -138,8 +147,8 @@ def _local_port_open(host: str, port: int, timeout: float = 1.0) -> bool: def _working_local_url(port: int) -> "str | None": - """Return a working loopback URL on this machine, or None if neither - 127.0.0.1 nor ::1 responds. Used as a fallback when external reachability fails.""" + """A working loopback URL on this machine, or None if neither 127.0.0.1 nor + ::1 responds. Fallback when external reachability fails.""" if _local_port_open("127.0.0.1", port): return f"http://127.0.0.1:{port}" if _local_port_open("::1", port): @@ -147,6 +156,51 @@ def _working_local_url(port: int) -> "str | None": return None +def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None": + """Return the IPv4 loopback URL when localhost won't reach 127.0.0.1. + + Local Studio binds to 127.0.0.1. Where localhost resolves to IPv6 only (::1), + http://localhost: fails (or hits a different process on ::1) even though + http://127.0.0.1: works. Return the IPv4 URL for the caller to surface. + """ + import socket + + if bind_host != "127.0.0.1" or not port or port <= 0: + return None + + ipv4_url = f"http://127.0.0.1:{port}" + + # Only warn once Studio is confirmed answering on IPv4 loopback. + if _working_local_url(port) != ipv4_url: + return None + + try: + addr_info = socket.getaddrinfo("localhost", port, socket.AF_UNSPEC, socket.SOCK_STREAM) + except Exception: + return None + + if not addr_info: + return None + + has_ipv4_loopback = False + has_ipv6_loopback = False + for family, _, _, _, sockaddr in addr_info: + if family == socket.AF_INET and sockaddr and sockaddr[0] == "127.0.0.1": + has_ipv4_loopback = True + elif family == socket.AF_INET6 and sockaddr: + host = sockaddr[0].split("%", 1)[0] + if host == "::1": + has_ipv6_loopback = True + + # A connection to ::1 is NOT evidence Studio is reachable there: Studio binds + # 127.0.0.1 only, so anything on ::1 is a different process. Dual-stack + # localhost is fine (browsers fall back to 127.0.0.1), so only the IPv6-only + # case strands the user. + if has_ipv6_loopback and not has_ipv4_loopback: + return ipv4_url + return None + + def _stdout_color_ok() -> bool: """Whether to emit ANSI color codes on stdout. Mirrors startup_banner.""" if os.environ.get("NO_COLOR", "").strip(): @@ -159,12 +213,28 @@ def _stdout_color_ok() -> bool: return False +def _print_localhost_ipv6_mismatch_warning(local_url: str, port: int) -> None: + """Warn that localhost points at ::1 while Studio is bound to 127.0.0.1.""" + use_color = _stdout_color_ok() + warn_c = "\033[38;5;215;1m" if use_color else "" + reset = "\033[0m" if use_color else "" + + print( + f"{warn_c} Warning: localhost resolves to IPv6 (::1), but Unsloth " + f"Studio is listening on 127.0.0.1 only. Open {local_url} instead of " + f"http://localhost:{port}.{reset}", + flush = True, + ) + + def _verify_global_reachability(display_host: str, port: int) -> None: """Probe check-host.net to confirm display_host:port is reachable from the - public internet. Synchronous so the caller can render output between the - banner URL section and the trailing stop hint. Bounded at ~15s; failures - are swallowed (the verifier failing is not Studio failing). Only meaningful - when bound to a wildcard host.""" + public internet. Synchronous so output lands between the banner URLs and the + stop hint. Bounded at ~15s; failures swallowed (verifier failing != Studio + failing). Only meaningful for a wildcard bind.""" + global _public_reachable + # Reset to "unknown" each run; set True/False only when the probe decides. + _public_reachable = None import ipaddress import json import time @@ -185,7 +255,7 @@ def _verify_global_reachability(display_host: str, port: int) -> None: url = f"http://{display_host}:{port}" - # Private / loopback / link-local addresses are not globally routable. + # Private/loopback/link-local addresses aren't globally routable. try: addr = ipaddress.ip_address(display_host) if addr.is_loopback or addr.is_private or addr.is_link_local: @@ -233,7 +303,7 @@ def _verify_global_reachability(display_host: str, port: int) -> None: continue if results and all(v is not None for v in results.values()): break - # Two decisive nodes is enough; stop polling early. + # Two decisive nodes is enough; stop early. decisive = [ v for v in results.values() @@ -257,12 +327,14 @@ def _verify_global_reachability(display_host: str, port: int) -> None: print("", flush = True) if ok_nodes: + _public_reachable = True print( f"{ok_c} Reachability check: {url}/ is reachable from the " f"public internet ({ok_nodes}/{total} probe nodes connected).{reset}", flush = True, ) elif err_nodes: + _public_reachable = False print( f"{err_c} Reachability check: {url}/ is NOT reachable from " f"the public internet ({err_nodes}/{total} probe nodes failed).{reset}", @@ -295,15 +367,14 @@ def _verify_global_reachability(display_host: str, port: int) -> None: flush = True, ) print( - f"{dim} ssh -L {port}:localhost:{port} " - f"@{display_host}{reset}", + f"{dim} ssh -L {port}:localhost:{port} " f"@{display_host}{reset}", flush = True, ) print( f"{dim} then open http://localhost:{port}/ in your browser.{reset}", flush = True, ) - # Only offer the local URL if loopback actually answers. + # Only offer the local URL if loopback answers. local_url = _working_local_url(port) if local_url: print( @@ -318,19 +389,66 @@ def _verify_global_reachability(display_host: str, port: int) -> None: flush = True, ) except urllib.error.URLError: - # Outbound HTTPS blocked; skip silently. + # Outbound HTTPS blocked; skip. pass except Exception: pass +def _emit_startup_output(host: str, port: int, display_host: str) -> None: + """Print the access banner plus any post-startup warnings. + + Extracted from ``_run`` so the banner/warning wiring is testable. The + ``localhost``-to-::1 mismatch warning and the wildcard reachability + check are mutually exclusive (the mismatch helper returns None for any + non-127.0.0.1 bind, and wildcard binds are never 127.0.0.1), so the + trailing stop hint is emitted exactly once. + """ + wildcard_bind = host in ("0.0.0.0", "::") + localhost_mismatch_url = _localhost_ipv6_mismatch_url(host, port) + # For wildcard binds, run the reachability check between the URL + # section and the stop hint so the stop hint stays last. + print_studio_access_banner( + port = port, + bind_host = host, + display_host = display_host, + include_stop_hint = not wildcard_bind and not localhost_mismatch_url, + ) + if localhost_mismatch_url: + _print_localhost_ipv6_mismatch_warning(localhost_mismatch_url, port) + print_studio_stop_hint() + elif wildcard_bind: + _verify_global_reachability(display_host, port) + _print_cloudflare_line() + print_studio_stop_hint() + + +def _print_cloudflare_line() -> None: + """Print the Cloudflare quick-tunnel URL for 0.0.0.0 binds, if one is up. + + Reads the module-level URL set by ``run_server``. Prints nothing when the + tunnel is disabled or failed -- failures are silently ignored. When the public + reachability probe just failed (``_public_reachable is False``) but the tunnel + is up, reword to point the user at the Cloudflare link as the way in. + """ + if not _cloudflare_url: + return + from startup_banner import stdout_supports_color + + accent = "\033[38;5;150;1m" + reset = "\033[0m" + if _public_reachable is False: + line = f" Use the secure link access via Cloudflare instead: {_cloudflare_url}" + else: + line = f" Secure link access via Cloudflare: {_cloudflare_url}" + print(f"{accent}{line}{reset}" if stdout_supports_color() else line) + + def _get_pid_on_port(port: int) -> "tuple[int, str] | None": - """Return (pid, process_name) of the process listening on *port*, or None. + """Return (pid, process_name) listening on *port*, or None. - Uses psutil when available. Falls back gracefully to None so callers - can still report the port conflict without process details. - - Works on Windows, macOS, and Linux wherever psutil is installed. + Uses psutil when available, else None so callers can still report the conflict + without process details. """ try: import psutil @@ -347,7 +465,7 @@ def _get_pid_on_port(port: int) -> "tuple[int, str] | None": except (psutil.NoSuchProcess, psutil.AccessDenied): return (conn.pid, "") except (psutil.AccessDenied, OSError) as e: - # psutil.net_connections() needs elevated privileges on some platforms + # net_connections() needs elevated privileges on some platforms. logger.debug("Failed to scan network connections for port %s: %s", port, e) return None @@ -355,19 +473,14 @@ def _get_pid_on_port(port: int) -> "tuple[int, str] | None": def _is_port_free(host: str, port: int) -> bool: """Check if a port is available for binding. - When *host* is ``0.0.0.0`` (wildcard), we also check whether anything - is already listening on ``127.0.0.1`` (and ``::1`` when IPv6 is - available). An SSH tunnel or similar process may hold the loopback - address while our wildcard bind still succeeds, making Unsloth Studio - unreachable via ``localhost``. - - Works on Windows, macOS, and Linux. + For a ``0.0.0.0`` wildcard host, also check whether anything is listening on + ``127.0.0.1`` (and ``::1`` when IPv6 exists): an SSH tunnel may hold loopback + while the wildcard bind succeeds, making Studio unreachable via ``localhost``. """ import socket - # 1. Can we bind to the requested address? - # Use getaddrinfo so both IPv4 ("0.0.0.0") and IPv6 ("::") hosts - # resolve to the correct address family automatically. + # 1. Can we bind to the requested address? getaddrinfo resolves both + # IPv4 and IPv6 to the right address family. try: addr_info = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM) family, socktype, proto, _, sockaddr = addr_info[0] @@ -377,9 +490,8 @@ def _is_port_free(host: str, port: int) -> bool: except OSError: return False - # 2. When binding to all interfaces, verify that localhost is not - # already claimed by another process (e.g. an SSH -L tunnel). - # We attempt a TCP connect -- if it succeeds something is listening. + # 2. On a wildcard bind, verify localhost isn't already claimed by another + # process (e.g. an SSH -L tunnel); a successful connect means it is. if host in ("0.0.0.0", "::"): for loopback, family in [ ("127.0.0.1", socket.AF_INET), @@ -389,24 +501,26 @@ def _is_port_free(host: str, port: int) -> bool: with socket.socket(family, socket.SOCK_STREAM) as s: s.settimeout(1) if s.connect_ex((loopback, port)) == 0: - # Connection succeeded -- port is taken on loopback + # Port is taken on loopback. return False except OSError: - # IPv6 disabled or other OS-level restriction -- skip + # IPv6 disabled or other OS-level restriction -- skip. continue return True -def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int: - """Find a free port starting from `start`, trying up to max_attempts ports.""" +def _find_free_port( + host: str, + start: int, + max_attempts: int = 20, +) -> int: + """Find a free port from `start`, trying up to max_attempts ports.""" for offset in range(max_attempts): candidate = start + offset if _is_port_free(host, candidate): return candidate - raise RuntimeError( - f"Could not find a free port in range {start}-{start + max_attempts - 1}" - ) + raise RuntimeError(f"Could not find a free port in range {start}-{start + max_attempts - 1}") from utils.paths.storage_roots import studio_root as _studio_root @@ -415,7 +529,7 @@ _PID_FILE = _studio_root() / "studio.pid" # Direct backend launches bypass the CLI's env re-export; do it here for # real custom roots so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR -# picks up the custom build. Skip for legacy-default to avoid flipping +# picks up the custom build. Skip legacy-default to avoid flipping # default-mode installs into env-override. try: _LEGACY_STUDIO_ROOT = (Path.home() / ".unsloth" / "studio").resolve() @@ -453,65 +567,77 @@ def _remove_pid_file(): def _graceful_shutdown(server = None): - """Explicitly shut down all subprocess backends and the uvicorn server. + """Shut down all subprocess backends and the uvicorn server. - Called from signal handlers to ensure child processes are cleaned up - before the parent exits. This is critical on Windows where atexit - handlers are unreliable after Ctrl+C. + Called from signal handlers to clean up children before exit. Critical on + Windows where atexit handlers are unreliable after Ctrl+C. """ _remove_pid_file() logger.info("Graceful shutdown initiated — cleaning up subprocesses...") - # 1. Shut down uvicorn server (releases the listening socket) + # 1. Shut down uvicorn (releases the listening socket). if server is not None: server.should_exit = True - # 2. Clean up inference subprocess (if instantiated) + # 2. Clean up inference subprocess (if instantiated). try: from core.inference.orchestrator import _inference_backend - if _inference_backend is not None: _inference_backend._shutdown_subprocess(timeout = 5.0) except Exception as e: logger.warning("Error shutting down inference subprocess: %s", e) - # 3. Clean up export subprocess (if instantiated) + # 3. Clean up export subprocess (if instantiated). try: from core.export.orchestrator import _export_backend - if _export_backend is not None: _export_backend._shutdown_subprocess(timeout = 5.0) except Exception as e: logger.warning("Error shutting down export subprocess: %s", e) - # 4. Clean up training subprocess (if active) + # 4. Clean up training subprocess (if active). try: from core.training.training import _training_backend - if _training_backend is not None: _training_backend.force_terminate() except Exception as e: logger.warning("Error shutting down training subprocess: %s", e) - # 5. Kill llama-server subprocess (if loaded) + # 5. Kill llama-server subprocess (if loaded). try: from routes.inference import _llama_cpp_backend - if _llama_cpp_backend is not None: _llama_cpp_backend._kill_process() except Exception as e: logger.warning("Error shutting down llama-server: %s", e) + # 6. Stop the Cloudflare tunnel (if started). + try: + from cloudflare_tunnel import stop_studio_tunnel + stop_studio_tunnel() + except Exception as e: + logger.warning("Error stopping Cloudflare tunnel: %s", e) + logger.info("All subprocesses cleaned up") # The uvicorn server instance -- set by run_server(), used by callers -# that need to tell the server to exit (e.g. signal handlers). +# that tell the server to exit (e.g. signal handlers). _server = None -# Shutdown event -- used to wake the main loop on signal +# Shutdown event -- wakes the main loop on signal. _shutdown_event = None +# trycloudflare.com URL for 0.0.0.0 binds (set by run_server, read by the banner); +# None when there is no tunnel (loopback, disabled, or a silently-ignored failure). +_cloudflare_url = None + +# Public reachability from the last _verify_global_reachability run, read by the +# Cloudflare banner line. True when the public ip:port probe confirmed reachable, +# False when it confirmed NOT reachable, None when the probe did not run or could +# not decide (timeout, blocked, private address). +_public_reachable = None + _DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist" @@ -519,9 +645,8 @@ _DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / " def _iter_frontend_fallback_candidates() -> "list[Path]": """Yield `studio/frontend/dist` paths to try when the default is missing. - Covers PATH-shadowed binaries whose __file__ resolves into a - site-packages tree that never received a vite build (e.g. plain - `pip install unsloth` from PyPI). + Covers PATH-shadowed binaries whose __file__ resolves into a site-packages + tree with no vite build (e.g. plain `pip install unsloth`). """ import ast import re @@ -547,20 +672,16 @@ def _iter_frontend_fallback_candidates() -> "list[Path]": src = finder.read_text(encoding = "utf-8") except OSError: continue - # Tolerate single- or multi-line dict literals; [^}]* still - # rejects nested dicts, which the setuptools template never - # emits for editable installs. - m = re.search( - r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S - ) + # Tolerate single/multi-line dict literals; [^}]* rejects nested + # dicts, which the setuptools editable template never emits. + m = re.search(r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S) if not m: continue try: mapping = ast.literal_eval(m.group(1)) except (SyntaxError, ValueError): continue - # Defensive: literal_eval can return a set / list / None if the - # matched literal is not a dict (regex captures `{...}`). + # literal_eval can return a set/list/None if `{...}` isn't a dict. if not isinstance(mapping, dict): continue studio_pkg = mapping.get("studio") @@ -570,10 +691,10 @@ def _iter_frontend_fallback_candidates() -> "list[Path]": def _resolve_frontend_path(frontend_path: Path) -> tuple[Optional[Path], list[Path]]: - """Pick a frontend dir that actually contains `index.html`. + """Pick a frontend dir that contains `index.html`. - Returns (chosen, attempted). `chosen` is None if nothing servable was - found; `attempted` is the full ordered list for diagnostics. + Returns (chosen, attempted). `chosen` is None if nothing servable was found; + `attempted` is the ordered list for diagnostics. """ attempted: list[Path] = [] seen: set[Path] = set() @@ -597,6 +718,94 @@ def _resolve_frontend_path(frontend_path: Path) -> tuple[Optional[Path], list[Pa return None, attempted +class _TeeStream: + """Mirror writes to the original stream and a session log file. + + Console behavior is unchanged (writes/returns delegate to the original + stream; Tauri's structured-stdout protocol and isatty probes see exactly + what they saw before). The file copy is best-effort: a full disk or a + closed handle must never break the console.""" + + def __init__(self, stream, log_fh): + self._stream = stream + self._log_fh = log_fh + + def write(self, data): + try: + self._log_fh.write(data) + except Exception: + pass + return self._stream.write(data) + + def flush(self): + try: + self._log_fh.flush() + except Exception: + pass + try: + self._stream.flush() + except Exception: + pass + + def __getattr__(self, name): + return getattr(self._stream, name) + + +def _setup_server_disk_logging(): + """Tee stdout/stderr to ~/.unsloth/studio/logs/server/ and aim + faulthandler at the same file so hard crashes (access violations / + SIGSEGV in the GPU runtime) leave a stack trace on disk. + + Also exports PYTHONFAULTHANDLER=1 so child Python processes (training + workers) dump native-crash stacks to their captured stderr. Keeps the + newest 20 session logs. Opt out with UNSLOTH_STUDIO_NO_FILE_LOG=1. + Returns the log path, or None when disabled/unavailable. + """ + if os.environ.get("UNSLOTH_STUDIO_NO_FILE_LOG") == "1": + return None + try: + from utils.paths import studio_root + log_dir = Path(studio_root()) / "logs" / "server" + except Exception: + home = ( + os.environ.get("UNSLOTH_STUDIO_HOME") + or os.environ.get("STUDIO_HOME") + or os.path.join(os.path.expanduser("~"), ".unsloth", "studio") + ) + log_dir = Path(home) / "logs" / "server" + try: + log_dir.mkdir(parents = True, exist_ok = True) + stamp = time.strftime("%Y%m%d-%H%M%S") + log_path = log_dir / f"server-{stamp}-pid{os.getpid()}.log" + # Line-buffered so the tail survives a hard kill; errors="replace" + # so a console encoding quirk can never take the server down. + log_fh = open(log_path, "w", encoding = "utf-8", errors = "replace", buffering = 1) + except Exception: + return None + + import faulthandler + + try: + faulthandler.enable(file = log_fh, all_threads = True) + except Exception: + pass + # Children (training workers) inherit: their native-crash stacks land on + # the stderr the server already captures. + os.environ.setdefault("PYTHONFAULTHANDLER", "1") + + sys.stdout = _TeeStream(sys.stdout, log_fh) + sys.stderr = _TeeStream(sys.stderr, log_fh) + + # Best-effort retention: keep the newest 20 session logs. + try: + logs = sorted(log_dir.glob("server-*.log"), key = lambda p: p.stat().st_mtime) + for old in logs[:-20]: + old.unlink(missing_ok = True) + except Exception: + pass + return log_path + + def run_server( host: str = "127.0.0.1", port: int = 8888, @@ -604,6 +813,7 @@ def run_server( silent: bool = False, api_only: bool = False, llama_parallel_slots: int = 1, + cloudflare: bool = True, ): """ Start the FastAPI server. @@ -613,25 +823,33 @@ def run_server( port: Port to bind to (auto-increments if in use) frontend_path: Path to frontend build directory (optional) silent: Suppress startup messages - api_only: Run API server only, no frontend serving (for Tauri desktop app) - llama_parallel_slots: Number of parallel slots for llama-server + api_only: API server only, no frontend (for Tauri desktop app) + llama_parallel_slots: parallel slots for llama-server Note: - Signal handlers are NOT registered here so that embedders - (e.g. Colab notebooks) keep their own interrupt semantics. - Standalone callers should register handlers after calling this. + Signal handlers are NOT registered here so embedders (e.g. Colab) keep + their own interrupt semantics; standalone callers register them after. """ global _server, _shutdown_event - # On Windows the default console encoding (cp1252) cannot encode emoji. - # Reconfigure stdout to UTF-8 so startup messages do not crash the server. + # Windows cp1252 can't encode emoji; reconfigure stdout to UTF-8. if sys.platform == "win32" and hasattr(sys.stdout, "reconfigure"): try: sys.stdout.reconfigure(encoding = "utf-8", errors = "replace") except Exception: pass - # Set env var BEFORE importing main so CORS middleware picks it up + # Persist a session log + native-crash stacks BEFORE importing main, so + # even import-time failures leave evidence on disk. Field report: Studio + # "terminates without a warning" -- a native crash in the GPU runtime + # kills the process with no Python traceback, and a desktop-shortcut + # console closes before anything can be read. Console-only logging made + # that undiagnosable. + _session_log = _setup_server_disk_logging() + if _session_log is not None and not silent: + print(f"Session log: {_session_log}") + + # Set env var BEFORE importing main so CORS middleware picks it up. if api_only: os.environ["UNSLOTH_API_ONLY"] = "1" @@ -643,13 +861,13 @@ def run_server( from threading import Thread, Event import uvicorn - from main import app, setup_frontend + from main import app, setup_frontend, _IS_COLAB from utils.paths import ensure_studio_directories - # Create all standard directories on startup + # Create all standard directories on startup. ensure_studio_directories() - # Auto-find free port if requested port is in use + # Auto-find a free port if the requested one is in use. if not _is_port_free(host, port): original_port = port blocker = _get_pid_on_port(port) @@ -659,9 +877,7 @@ def run_server( print("=" * 50) if blocker: pid, name = blocker - print( - f"Port {original_port} is already in use by " f"{name} (PID {pid})." - ) + print(f"Port {original_port} is already in use by " f"{name} (PID {pid}).") else: print(f"Port {original_port} is already in use.") print(f"Unsloth Studio will use port {port} instead.") @@ -669,14 +885,13 @@ def run_server( print("=" * 50) print("") - # Setup frontend if path provided (skip in api-only mode). - # Falls back through alternate locations if the default lacks a built - # dist; errors out loudly rather than silently serving 404 on `/`. + # Setup frontend (skip in api-only). Falls back through alternate locations if + # the default lacks a built dist; errors loudly rather than 404 on `/`. if frontend_path and not api_only: chosen, attempted = _resolve_frontend_path(Path(frontend_path)) if chosen is not None and setup_frontend(app, chosen): if not silent: - # Resolve so logs always show an absolute path for support. + # Resolve so logs show an absolute path for support. try: display = chosen.resolve() except OSError: @@ -688,9 +903,8 @@ def run_server( or os.environ.get("STUDIO_HOME") or str(Path.home() / ".unsloth" / "studio") ) - # Windows ships the user-facing shim at $STUDIO_HOME/bin/unsloth.exe - # (a hardlink to the venv exe); Linux/macOS use the venv binary - # at $STUDIO_HOME/unsloth_studio/bin/unsloth. + # Windows shim: $STUDIO_HOME/bin/unsloth.exe; Linux/macOS venv binary: + # $STUDIO_HOME/unsloth_studio/bin/unsloth. home = Path(home_str).expanduser() if sys.platform == "win32": installer_bin = home / "bin" / "unsloth.exe" @@ -712,7 +926,7 @@ def run_server( " - reinstall: curl -fsSL https://unsloth.ai/install.sh | sh" ) - # Resolve once; shared by the log rewrite and the banner. + # Resolve once; shared by the log rewrite and banner. display_host = _resolve_external_ip() if host == "0.0.0.0" else host _install_uvicorn_startup_log_rewrite(host, display_host) @@ -727,28 +941,32 @@ def run_server( ready_event.set() # server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own. - config = uvicorn.Config( - app, + config_kwargs = dict( host = host, port = port, log_level = "info", access_log = False, server_header = False, ) + # Colab only: trust X-Forwarded-* from Colab's reverse proxy so the app sees + # the real https origin. forwarded_allow_ips="*" is safe in Colab's + # single-user sandbox but too lax for local/standalone, so leave uvicorn's + # loopback-only default elsewhere. + if _IS_COLAB: + config_kwargs["proxy_headers"] = True + config_kwargs["forwarded_allow_ips"] = "*" + config = uvicorn.Config(app, **config_kwargs) _server = _ReadyServer(config) _shutdown_event = Event() - # Expose the actual bound port so request-handling code can build - # loopback URLs that point at the real backend, not whatever port a - # reverse proxy or tunnel exposed in the request URL. Only publish - # an explicit value when we know the concrete port; for ephemeral - # binds (port==0) leave it unset and let request handlers fall back - # to the ASGI request scope or request.base_url. + # Expose the actual bound port so handlers build loopback URLs at the real + # backend, not whatever a proxy/tunnel exposed. For ephemeral binds (port==0) + # leave it unset so handlers fall back to the request scope / base_url. app.state.server_port = port if port and port > 0 else None app.state.llama_parallel_slots = llama_parallel_slots - # Expose a shutdown callable via app.state before the server can accept - # requests so /api/shutdown is available as soon as readiness is published. + # Expose a shutdown callable before the server accepts requests so + # /api/shutdown is ready as soon as readiness publishes. def _trigger_shutdown(): _graceful_shutdown(_server) if _shutdown_event is not None: @@ -756,23 +974,27 @@ def run_server( app.state.trigger_shutdown = _trigger_shutdown - # Run server in a daemon thread + # Run server in a daemon thread with explicit new_event_loop() + + # run_until_complete() (not asyncio.run) so nest_asyncio's patches don't + # interfere when Colab/IPython already runs a loop on the main thread. def _run(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) try: - asyncio.run(_server.serve()) + loop.run_until_complete(_server.serve()) except BaseException as exc: startup_errors.append(exc) startup_failed.set() finally: + loop.close() if not ready_event.is_set(): startup_failed.set() thread = Thread(target = _run, daemon = True) thread.start() - # Wait until uvicorn has completed lifespan startup and bound sockets, or - # until the server exits/fails before startup. This intentionally has no - # correctness deadline: a slow but live startup should remain in progress. + # Wait until uvicorn finishes lifespan startup and binds sockets, or until it + # exits/fails first. No deadline: a slow but live startup stays in progress. try: while not ready_event.is_set(): if startup_failed.is_set() or not thread.is_alive(): @@ -792,35 +1014,38 @@ def run_server( atexit.register(_remove_pid_file) - # Output port for Tauri to parse when in api-only mode. Emit only after - # uvicorn sockets are bound and FastAPI lifespan/startup has completed. + # Output port for Tauri (api-only), only after sockets bind and startup done. if api_only: print(f"TAURI_PORT={port}", flush = True) + # Free trycloudflare.com tunnel for 0.0.0.0 binds (the raw ip:port is often + # unreachable). Started pre-banner and even when silent so the CLI banner can + # read app.state.cloudflare_url; torn down by _graceful_shutdown. + global _cloudflare_url + _cloudflare_url = None + app.state.cloudflare_url = None + _cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB + if _cloudflare_enabled: + try: # best-effort: any failure must not block startup + from cloudflare_tunnel import start_studio_tunnel + _cloudflare_url = start_studio_tunnel(port) + app.state.cloudflare_url = _cloudflare_url + except Exception as e: + logger.debug("Cloudflare tunnel skipped: %s", e) + if not silent: - wildcard_bind = host in ("0.0.0.0", "::") - # For wildcard binds, run the reachability check between the URL - # section and the stop hint so the stop hint stays last on screen. - print_studio_access_banner( - port = port, - bind_host = host, - display_host = display_host, - include_stop_hint = not wildcard_bind, - ) - if wildcard_bind: - _verify_global_reachability(display_host, port) - print_studio_stop_hint() + _emit_startup_output(host, port, display_host) return app -# For direct execution (also invoked by CLI via os.execvp / subprocess) +# For direct execution (also invoked by CLI via os.execvp / subprocess). if __name__ == "__main__": import argparse import signal import traceback - # Ensure stderr can handle Unicode on Windows (tracebacks with non-ASCII paths) + # Ensure stderr handles Unicode on Windows (non-ASCII path tracebacks). if sys.platform == "win32" and hasattr(sys.stderr, "reconfigure"): try: sys.stderr.reconfigure(encoding = "utf-8", errors = "replace") @@ -846,9 +1071,15 @@ if __name__ == "__main__": action = "store_true", help = "API server only, no frontend (for Tauri)", ) - # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 - # applies only to direct backend launches; `unsloth studio run` - # always passes its own value (4) explicitly. + parser.add_argument( + "--cloudflare", + action = argparse.BooleanOptionalAction, + default = True, + help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 " + "(default on; --no-cloudflare to disable)", + ) + # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct + # backend launches; `unsloth studio run` always passes its own value (4). _PARALLEL_MIN = 1 _PARALLEL_MAX = 64 _PARALLEL_DEFAULT_PLAIN = 1 @@ -873,6 +1104,7 @@ if __name__ == "__main__": silent = args.silent, api_only = args.api_only, llama_parallel_slots = args.parallel, + cloudflare = args.cloudflare, ) if args.frontend is not None: kwargs["frontend_path"] = Path(args.frontend) @@ -886,13 +1118,11 @@ if __name__ == "__main__": sys.stderr.write("=" * 60 + "\n") traceback.print_exc(file = sys.stderr) sys.stderr.write("\n") - sys.stderr.write( - "If a package is missing, try re-running: unsloth studio setup\n" - ) + sys.stderr.write("If a package is missing, try re-running: unsloth studio setup\n") sys.stderr.flush() sys.exit(1) - # Signal handler -- ensures subprocess cleanup on Ctrl+C + # Signal handler -- ensures subprocess cleanup on Ctrl+C. def _signal_handler(signum, frame): _graceful_shutdown(_server) _shutdown_event.set() @@ -900,13 +1130,12 @@ if __name__ == "__main__": signal.signal(signal.SIGINT, _signal_handler) signal.signal(signal.SIGTERM, _signal_handler) - # On Windows, some terminals send SIGBREAK for Ctrl+C / Ctrl+Break + # On Windows, some terminals send SIGBREAK for Ctrl+C / Ctrl+Break. if hasattr(signal, "SIGBREAK"): signal.signal(signal.SIGBREAK, _signal_handler) - # Keep running until shutdown signal. - # NOTE: Event.wait() without a timeout blocks at the C level on Linux, - # which prevents Python from delivering SIGINT (Ctrl+C). Using a - # short timeout in a loop lets the interpreter process pending signals. + # Keep running until shutdown signal. Event.wait() without a timeout blocks at + # the C level on Linux, preventing SIGINT delivery; a short timeout in a loop + # lets the interpreter process pending signals. while not _shutdown_event.is_set(): _shutdown_event.wait(timeout = 1) diff --git a/studio/backend/startup_banner.py b/studio/backend/startup_banner.py index 2bda4357ba..7bfca61b73 100644 --- a/studio/backend/startup_banner.py +++ b/studio/backend/startup_banner.py @@ -3,7 +3,7 @@ """Terminal banner for Studio startup. -Stdlib only — safe to import without the rest of the backend (no structlog/uvicorn). +Stdlib only — safe to import without the rest of the backend. """ from __future__ import annotations @@ -34,7 +34,7 @@ def print_port_in_use_notice(original_port: int, new_port: int) -> None: def print_studio_stop_hint() -> None: - """Print the trailing stop hint + closing divider. Separate from the main + """Print the trailing stop hint + closing divider, separate from the banner so callers can interleave content (e.g. a reachability check).""" use_color = stdout_supports_color() dim = "\033[38;5;245m" @@ -67,7 +67,7 @@ def print_studio_access_banner( display_host: str, include_stop_hint: bool = True, ) -> None: - """Pretty-print URLs after the server is listening. Set + """Pretty-print URLs once the server is listening. Set ``include_stop_hint=False`` to omit the trailing stop block; pair with :func:`print_studio_stop_hint` after inserting your own content.""" use_color = stdout_supports_color() @@ -96,8 +96,8 @@ def print_studio_access_banner( listen_all = bind_host in ("0.0.0.0", "::") loopback_bind = bind_host in ("127.0.0.1", "localhost", "::1") - # Use loopback URL only when the server is reachable on loopback; - # otherwise show the actual bound address. + # Use the loopback URL only when reachable on loopback; otherwise show + # the actual bound address. primary_url = loopback_url if listen_all or loopback_bind else external_url tip_url = alt_local if listen_all or loopback_bind else external_url api_base = primary_url diff --git a/studio/backend/state/tool_policy.py b/studio/backend/state/tool_policy.py index 9343a39806..9b0fc7d6cb 100644 --- a/studio/backend/state/tool_policy.py +++ b/studio/backend/state/tool_policy.py @@ -21,9 +21,7 @@ def get_tool_policy() -> Optional[bool]: def set_tool_policy(value: Optional[bool]) -> None: if value is not None and not isinstance(value, bool): - raise TypeError( - f"tool_policy must be Optional[bool], got {type(value).__name__}" - ) + raise TypeError(f"tool_policy must be Optional[bool], got {type(value).__name__}") global _tool_policy _tool_policy = value diff --git a/studio/backend/storage/mcp_servers_db.py b/studio/backend/storage/mcp_servers_db.py new file mode 100644 index 0000000000..6482bae140 --- /dev/null +++ b/studio/backend/storage/mcp_servers_db.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import sqlite3 +import threading +from datetime import datetime, timezone +from typing import Optional + +from utils.paths import studio_db_path, ensure_dir + +_schema_lock = threading.Lock() +_schema_ready = False + + +def _ensure_schema(conn: sqlite3.Connection) -> None: + conn.execute("PRAGMA journal_mode=WAL") + conn.execute( + """ + CREATE TABLE IF NOT EXISTS mcp_servers ( + id TEXT NOT NULL PRIMARY KEY, + display_name TEXT NOT NULL, + url TEXT NOT NULL, + headers_json TEXT, + is_enabled INTEGER NOT NULL DEFAULT 1, + use_oauth INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ) + """ + ) + # Backfill use_oauth for pre-existing DBs. + cols = {r["name"] for r in conn.execute("PRAGMA table_info(mcp_servers)").fetchall()} + if "use_oauth" not in cols: + conn.execute("ALTER TABLE mcp_servers ADD COLUMN use_oauth INTEGER NOT NULL DEFAULT 0") + + +def get_connection() -> sqlite3.Connection: + global _schema_ready + db_path = studio_db_path() + ensure_dir(db_path.parent) + conn = sqlite3.connect(str(db_path)) + conn.row_factory = sqlite3.Row + if not _schema_ready: + with _schema_lock: + if not _schema_ready: + try: + _ensure_schema(conn) + _schema_ready = True + except Exception: + conn.close() + raise + return conn + + +def create_server( + id: str, + display_name: str, + url: str, + headers_json: Optional[str] = None, + is_enabled: bool = True, + use_oauth: bool = False, +) -> None: + now = datetime.now(timezone.utc).isoformat() + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO mcp_servers + (id, display_name, url, headers_json, + is_enabled, use_oauth, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + id, + display_name, + url, + headers_json, + int(is_enabled), + int(use_oauth), + now, + now, + ), + ) + conn.commit() + finally: + conn.close() + + +def update_server(id: str, changes: dict) -> bool: + """Apply column updates and bump ``updated_at``. Returns True on a hit.""" + if not changes: + return False + bool_cols = {"is_enabled", "use_oauth"} + sets, params = [], [] + for col, value in changes.items(): + sets.append(f"{col} = ?") + params.append(int(value) if col in bool_cols else value) + sets.append("updated_at = ?") + params.extend([datetime.now(timezone.utc).isoformat(), id]) + + conn = get_connection() + try: + cursor = conn.execute( + f"UPDATE mcp_servers SET {', '.join(sets)} WHERE id = ?", + params, + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + +def delete_server(id: str) -> bool: + conn = get_connection() + try: + cursor = conn.execute("DELETE FROM mcp_servers WHERE id = ?", (id,)) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + +def get_server(id: str) -> Optional[dict]: + conn = get_connection() + try: + row = conn.execute("SELECT * FROM mcp_servers WHERE id = ?", (id,)).fetchone() + return dict(row) if row else None + finally: + conn.close() + + +def list_servers() -> list[dict]: + conn = get_connection() + try: + rows = conn.execute("SELECT * FROM mcp_servers ORDER BY created_at").fetchall() + return [dict(row) for row in rows] + finally: + conn.close() diff --git a/studio/backend/storage/providers_db.py b/studio/backend/storage/providers_db.py index ca47fcbd80..07165cbe70 100644 --- a/studio/backend/storage/providers_db.py +++ b/studio/backend/storage/providers_db.py @@ -1,14 +1,11 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -SQLite storage for external LLM provider configurations. +"""SQLite storage for external LLM provider configurations. -Follows the same pattern as studio_db.py — module-level functions, -raw sqlite3, WAL mode, per-function connections. - -NOTE: API keys are NOT stored here. They live only in the browser -(localStorage) and are sent encrypted per-request. +Same pattern as studio_db.py (module-level functions, raw sqlite3, WAL, +per-function connections). API keys are NOT stored here: they live only in +the browser (localStorage) and are sent encrypted per-request. """ import logging @@ -26,7 +23,7 @@ _schema_ready = False def _ensure_schema(conn: sqlite3.Connection) -> None: - """Create the llm_providers table if it doesn't exist. Called once per process.""" + """Create the llm_providers table if absent. Called once per process.""" conn.execute("PRAGMA journal_mode=WAL") conn.execute( """ @@ -62,12 +59,7 @@ def get_connection() -> sqlite3.Connection: return conn -def create_provider( - id: str, - provider_type: str, - display_name: str, - base_url: str, -) -> None: +def create_provider(id: str, provider_type: str, display_name: str, base_url: str) -> None: """Insert a new provider configuration.""" now = datetime.now(timezone.utc).isoformat() conn = get_connection() @@ -145,9 +137,7 @@ def list_providers() -> list[dict]: """List all provider configurations, ordered by creation time.""" conn = get_connection() try: - rows = conn.execute( - "SELECT * FROM llm_providers ORDER BY created_at" - ).fetchall() + rows = conn.execute("SELECT * FROM llm_providers ORDER BY created_at").fetchall() return [dict(row) for row in rows] finally: conn.close() diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py new file mode 100644 index 0000000000..601fb73d0a --- /dev/null +++ b/studio/backend/storage/rag_db.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""SQLite storage for the RAG engine. + +Same pattern as providers_db.py / studio_db.py (module functions, raw sqlite3, +WAL, per-call connections, lazy schema), but every connection also loads +sqlite-vec (vec0 needs it per-connection). If it cannot load, RAG_AVAILABLE is +False and get_connection() raises rather than failing import. + +One rag.db holds the ``documents`` / ``chunks`` model, the FTS5 lexical index +(``chunks_fts``) and the sqlite-vec dense index (``chunks_vec``, created lazily +by ensure_vec once the embedding dim is known, since vec0 bakes the dim into the +column type). +""" + +import logging +import sqlite3 +import threading + +logger = logging.getLogger(__name__) + +from utils.paths import rag_db_path, ensure_dir + +# Optional dep: import must never crash this module (imported unconditionally). +try: + import sqlite_vec + RAG_AVAILABLE = True +except Exception as exc: # noqa: BLE001 - any import failure disables RAG + sqlite_vec = None + RAG_AVAILABLE = False + logger.warning("RAG unavailable: sqlite-vec could not be imported (%s)", exc) + +_RAG_UNAVAILABLE_MSG = "RAG unavailable: sqlite-vec extension could not be loaded" + +_schema_lock = threading.Lock() +_schema_ready = False + + +def _ensure_schema(conn: sqlite3.Connection) -> None: + """Create the RAG tables if absent (once per process). ``chunks_vec`` is + skipped: its column type needs the embedding dim, so ensure_vec() makes it + lazily at first ingest.""" + conn.execute("PRAGMA journal_mode=WAL") + conn.executescript( + """ + CREATE TABLE IF NOT EXISTS knowledge_bases ( + id TEXT NOT NULL PRIMARY KEY, + name TEXT NOT NULL, + description TEXT, + embedding_model TEXT, + created_at TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS documents ( + id TEXT NOT NULL PRIMARY KEY, + scope TEXT NOT NULL, + kb_id TEXT, + thread_id TEXT, + filename TEXT NOT NULL, + sha256 TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + error TEXT, + num_chunks INTEGER NOT NULL DEFAULT 0, + stored_path TEXT, + created_at TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_documents_scope ON documents(scope); + CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(scope, sha256); + + CREATE TABLE IF NOT EXISTS chunks ( + id TEXT NOT NULL PRIMARY KEY, + document_id TEXT NOT NULL, + scope TEXT NOT NULL, + chunk_index INTEGER NOT NULL, + text TEXT NOT NULL, + page_number INTEGER, + source_page_index INTEGER, + token_count INTEGER, + kind TEXT NOT NULL DEFAULT 'text', + pdf_regions_json TEXT + ); + CREATE INDEX IF NOT EXISTS idx_chunks_scope ON chunks(scope); + CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(document_id); + + CREATE TABLE IF NOT EXISTS ingestion_jobs ( + id TEXT NOT NULL PRIMARY KEY, + document_id TEXT NOT NULL, + scope TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + stage TEXT, + progress REAL NOT NULL DEFAULT 0.0, + error TEXT, + created_at TEXT NOT NULL + ); + + CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5( + text, + chunk_id UNINDEXED, + scope UNINDEXED, + tokenize='porter unicode61' + ); + """ + ) + + +def get_connection() -> sqlite3.Connection: + """Open rag.db (WAL + sqlite-vec loaded, schema created once). Raises if the extension is unavailable.""" + global _schema_ready + if not RAG_AVAILABLE: + raise RuntimeError(_RAG_UNAVAILABLE_MSG) + + db_path = rag_db_path() + ensure_dir(db_path.parent) + conn = sqlite3.connect(str(db_path)) + conn.row_factory = sqlite3.Row + try: + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + except Exception as exc: # noqa: BLE001 + conn.close() + raise RuntimeError(_RAG_UNAVAILABLE_MSG) from exc + + if not _schema_ready: + with _schema_lock: + if not _schema_ready: + try: + _ensure_schema(conn) + _schema_ready = True + except Exception: + conn.close() + raise + return conn + + +def ensure_vec(conn: sqlite3.Connection, dim: int) -> None: + """Create the dense ``chunks_vec`` table once the embedding dim is known + (vec0 bakes it into the column type). Idempotent; dim fixed per db.""" + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(" + f"scope TEXT partition key, " + f"chunk_id TEXT, " + f"embedding float[{int(dim)}] distance_metric=cosine)" + ) + + +def vec_table_exists(conn: sqlite3.Connection) -> bool: + """True if the dense ``chunks_vec`` table exists.""" + row = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name='chunks_vec'" + ).fetchone() + return row is not None diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index de89b6cbd2..da8d9b5e66 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -1,28 +1,28 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -SQLite storage for training run history and metrics. +"""SQLite storage for training run history and metrics. -Follows the same pattern as auth/storage.py — module-level functions, -raw sqlite3, per-function connections. Enhancements over auth: - - WAL mode for concurrent read/write access - - PRAGMA foreign_keys = ON for CASCADE deletes +Like auth/storage.py (module-level functions, raw sqlite3, per-function +connections) plus WAL mode and PRAGMA foreign_keys = ON for CASCADE deletes. """ import json import logging import os import platform +import re +import shutil import sqlite3 import threading from datetime import datetime, timezone +from pathlib import Path logger = logging.getLogger(__name__) from typing import Any, Iterable, Optional -from utils.paths import studio_db_path, ensure_dir +from utils.paths import project_workspaces_root, studio_db_path, ensure_dir def _denied_path_prefixes() -> list[str]: @@ -31,8 +31,7 @@ def _denied_path_prefixes() -> list[str]: if system == "Linux": return ["/proc", "/sys", "/dev", "/etc", "/boot", "/run"] if system == "Darwin": - # realpath() resolves /etc -> /private/etc, /tmp -> /private/tmp on macOS, - # so include the /private variants to avoid bypasses. + # macOS realpath() resolves /etc -> /private/etc etc; include /private variants. return [ "/System", "/Library", @@ -55,6 +54,75 @@ def _denied_path_prefixes() -> list[str]: _schema_lock = threading.Lock() _schema_ready = False _SQLITE_IN_CHUNK_SIZE = 900 +_PROJECT_WORKSPACE_SUBDIRS = ("sandbox",) + + +def _project_slug(name: str) -> str: + slug = re.sub(r"[^A-Za-z0-9._-]+", "-", name.strip()).strip(".-_") + return slug[:48] or "project" + + +def _default_project_root(project: dict) -> str: + project_id = str(project["id"]) + suffix = re.sub(r"[^A-Za-z0-9_-]+", "-", project_id)[:8].strip("-_") or "project" + folder_name = f"{_project_slug(str(project.get('name') or 'Project'))}-{suffix}" + return str(project_workspaces_root() / folder_name) + + +def _ensure_project_workspace(root_path: str) -> str: + root = Path(root_path).expanduser() + root_resolved = ensure_dir(root).resolve() + for subdir in _PROJECT_WORKSPACE_SUBDIRS: + ensure_dir(root_resolved / subdir) + return str(root_resolved) + + +def _delete_project_workspace(project: dict) -> None: + root_path = project.get("rootPath") + if not root_path: + return + root = Path(root_path).expanduser() + try: + root_resolved = root.resolve(strict = False) + except (OSError, RuntimeError, ValueError): + logger.warning("Skipping project workspace delete for invalid path %r", root_path) + return + + project_id = str(project["id"]) + suffix = re.sub(r"[^A-Za-z0-9_-]+", "-", project_id)[:8].strip("-_") or "project" + if not root_resolved.name.endswith(f"-{suffix}"): + logger.warning( + "Skipping project workspace delete for unexpected project path %s", + root_resolved, + ) + return + if root_resolved.parent == root_resolved or root_resolved == Path.home().resolve(): + logger.warning( + "Skipping project workspace delete for unsafe project path %s", + root_resolved, + ) + return + check = ( + os.path.normcase(str(root_resolved)) + if platform.system() == "Windows" + else str(root_resolved) + ) + for prefix in _denied_path_prefixes(): + if check == prefix or check.startswith(prefix + os.sep): + logger.warning( + "Skipping project workspace delete under denied path %s", + root_resolved, + ) + return + if not root_resolved.exists(): + return + if root_resolved.is_symlink() or not root_resolved.is_dir(): + logger.warning( + "Skipping project workspace delete for non-directory path %s", + root_resolved, + ) + return + shutil.rmtree(root_resolved) def _ensure_schema(conn: sqlite3.Connection) -> None: @@ -81,9 +149,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) - existing_cols = { - row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall() - } + existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()} if "display_name" not in existing_cols: conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT") conn.execute( @@ -103,12 +169,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)" - ) - # Use COLLATE NOCASE on Windows so C:\Models and c:\models dedup via the - # UNIQUE constraint. On Linux/macOS (case-sensitive FS) keep the default - # BINARY collation so /Models and /models remain distinct. + conn.execute("CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)") + # Windows: COLLATE NOCASE so C:\Models and c:\models dedup. Elsewhere keep + # case-sensitive BINARY so /Models and /models stay distinct. collation = "COLLATE NOCASE" if platform.system() == "Windows" else "" conn.execute( f""" @@ -119,6 +182,27 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_projects ( + id TEXT NOT NULL PRIMARY KEY, + name TEXT NOT NULL, + instructions TEXT, + root_path TEXT, + archived INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ) + """ + ) + chat_project_cols = { + row[1] for row in conn.execute("PRAGMA table_info(chat_projects)").fetchall() + } + if "root_path" not in chat_project_cols: + conn.execute("ALTER TABLE chat_projects ADD COLUMN root_path TEXT") + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_projects_archived_updated_at ON chat_projects(archived, updated_at)" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS chat_threads ( @@ -127,24 +211,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: model_type TEXT NOT NULL, model_id TEXT, pair_id TEXT, + project_id TEXT, archived INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, openai_code_exec_container_id TEXT, - anthropic_code_exec_container_id TEXT + anthropic_code_exec_container_id TEXT, + FOREIGN KEY(project_id) REFERENCES chat_projects(id) ON DELETE CASCADE ) """ ) chat_thread_cols = { row[1] for row in conn.execute("PRAGMA table_info(chat_threads)").fetchall() } + if "project_id" not in chat_thread_cols: + conn.execute("ALTER TABLE chat_threads ADD COLUMN project_id TEXT") if "openai_code_exec_container_id" not in chat_thread_cols: - conn.execute( - "ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT" - ) + conn.execute("ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT") if "anthropic_code_exec_container_id" not in chat_thread_cols: - conn.execute( - "ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT" - ) + conn.execute("ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT") conn.execute( """ CREATE TABLE IF NOT EXISTS chat_messages ( @@ -162,8 +246,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute( "CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)" ) + conn.execute("CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)") conn.execute( - "CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)" + "CREATE INDEX IF NOT EXISTS idx_chat_threads_project_id ON chat_threads(project_id)" ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_chat_messages_thread_id_created_at ON chat_messages(thread_id, created_at)" @@ -177,6 +262,15 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS app_settings ( + key TEXT NOT NULL PRIMARY KEY, + value_json TEXT NOT NULL, + updated_at TEXT NOT NULL + ) + """ + ) conn.execute( """ CREATE TABLE IF NOT EXISTS chat_settings_quarantine ( @@ -188,15 +282,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) - # Server-side import ledger so a studio.db wipe correctly re-triggers - # the legacy Dexie import. The previous boolean localStorage sentinel - # (`unsloth_chat_legacy_imported_to_studio_db`) is non-recoverable: - # if studio.db is recreated while the browser keeps the flag, legacy - # Dexie threads are silently hidden from the sidebar. The ledger - # lives inside studio.db so it disappears together with the data it - # is supposed to track, which is the recovery the boolean lacked. - # Keyed by legacy thread id; per-thread is sufficient because Dexie - # is read-only after this PR (a thread's message set does not grow). + # Import ledger inside studio.db (vs. a localStorage boolean) so a db wipe + # re-triggers the legacy Dexie import instead of silently hiding threads. + # Keyed by legacy thread id; Dexie is read-only so per-thread suffices. conn.execute( """ CREATE TABLE IF NOT EXISTS chat_legacy_imports ( @@ -205,6 +293,195 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) WITHOUT ROWID """ ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS prompt_entries ( + id TEXT NOT NULL PRIMARY KEY, + name TEXT NOT NULL, + text TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_prompt_entries_created_at ON prompt_entries(created_at)" + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS prompt_lists ( + id TEXT NOT NULL PRIMARY KEY, + name TEXT NOT NULL, + items_json TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_prompt_lists_created_at ON prompt_lists(created_at)" + ) + + +def _prompt_entry_from_row(row: sqlite3.Row) -> dict: + return { + "id": row["id"], + "name": row["name"], + "text": row["text"], + "createdAt": row["created_at"], + "updatedAt": row["updated_at"], + } + + +def list_prompt_entries() -> list[dict]: + conn = get_connection() + try: + rows = conn.execute("SELECT * FROM prompt_entries ORDER BY created_at DESC").fetchall() + return [_prompt_entry_from_row(r) for r in rows] + finally: + conn.close() + + +def upsert_prompt_entry(entry: dict) -> dict: + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO prompt_entries (id, name, text, created_at, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + name = excluded.name, + text = excluded.text, + updated_at = excluded.updated_at + """, + ( + entry["id"], + entry["name"], + entry["text"], + entry["createdAt"], + entry["updatedAt"], + ), + ) + conn.commit() + return entry + finally: + conn.close() + + +def delete_prompt_entry(entry_id: str) -> None: + conn = get_connection() + try: + conn.execute("DELETE FROM prompt_entries WHERE id = ?", (entry_id,)) + conn.commit() + finally: + conn.close() + + +def bulk_upsert_prompt_entries(entries: list[dict]) -> int: + if not entries: + return 0 + conn = get_connection() + try: + conn.executemany( + """ + INSERT INTO prompt_entries (id, name, text, created_at, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + name = excluded.name, + text = excluded.text, + updated_at = excluded.updated_at + """, + [(e["id"], e["name"], e["text"], e["createdAt"], e["updatedAt"]) for e in entries], + ) + conn.commit() + return len(entries) + finally: + conn.close() + + +def _prompt_list_from_row(row: sqlite3.Row) -> dict: + return { + "id": row["id"], + "name": row["name"], + "items": json.loads(row["items_json"]), + "createdAt": row["created_at"], + "updatedAt": row["updated_at"], + } + + +def list_prompt_lists_db() -> list[dict]: + conn = get_connection() + try: + rows = conn.execute("SELECT * FROM prompt_lists ORDER BY created_at DESC").fetchall() + return [_prompt_list_from_row(r) for r in rows] + finally: + conn.close() + + +def upsert_prompt_list(lst: dict) -> dict: + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO prompt_lists (id, name, items_json, created_at, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + name = excluded.name, + items_json = excluded.items_json, + updated_at = excluded.updated_at + """, + ( + lst["id"], + lst["name"], + json.dumps(lst["items"]), + lst["createdAt"], + lst["updatedAt"], + ), + ) + conn.commit() + return lst + finally: + conn.close() + + +def delete_prompt_list_db(list_id: str) -> None: + conn = get_connection() + try: + conn.execute("DELETE FROM prompt_lists WHERE id = ?", (list_id,)) + conn.commit() + finally: + conn.close() + + +def bulk_upsert_prompt_lists(lists: list[dict]) -> int: + if not lists: + return 0 + conn = get_connection() + try: + conn.executemany( + """ + INSERT INTO prompt_lists (id, name, items_json, created_at, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + name = excluded.name, + items_json = excluded.items_json, + updated_at = excluded.updated_at + """, + [ + ( + lst["id"], + lst["name"], + json.dumps(lst["items"]), + lst["createdAt"], + lst["updatedAt"], + ) + for lst in lists + ], + ) + conn.commit() + return len(lists) + finally: + conn.close() def get_connection() -> sqlite3.Connection: @@ -214,7 +491,7 @@ def get_connection() -> sqlite3.Connection: ensure_dir(db_path.parent) conn = sqlite3.connect(str(db_path)) conn.row_factory = sqlite3.Row - # foreign_keys is session-scoped, must be set per connection + # foreign_keys is session-scoped; set per connection conn.execute("PRAGMA foreign_keys=ON") if not _schema_ready: with _schema_lock: @@ -402,9 +679,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: try: run["loss_sparkline"] = json.loads(sparkline) except (json.JSONDecodeError, TypeError): - logger.debug( - "Failed to parse loss_sparkline for run %s", run.get("id") - ) + logger.debug("Failed to parse loss_sparkline for run %s", run.get("id")) run["loss_sparkline"] = None runs.append(run) return {"runs": runs, "total": total} @@ -480,9 +755,7 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]: try: run["loss_sparkline"] = json.loads(sparkline) except (json.JSONDecodeError, TypeError): - logger.debug( - "Failed to parse loss_sparkline for output_dir %s", output_dir - ) + logger.debug("Failed to parse loss_sparkline for output_dir %s", output_dir) run["loss_sparkline"] = None return run finally: @@ -606,9 +879,8 @@ def add_scan_folder(path: str) -> dict: if not os.access(normalized, os.R_OK | os.X_OK): raise ValueError("Path is not readable") - # On Windows, use normcase for denylist comparison but store the - # original-cased path so downstream consumers see the native - # drive-letter casing the user expects (e.g. C:\Models, not c:\models). + # Windows: normcase for the denylist check but store original casing + # so consumers see the native drive-letter casing (e.g. C:\Models). is_win = platform.system() == "Windows" check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): @@ -618,8 +890,7 @@ def add_scan_folder(path: str) -> dict: conn = get_connection() try: now = datetime.now(timezone.utc).isoformat() - # On Windows, use case-insensitive lookup so C:\Models and c:\models - # dedup correctly while preserving the originally-stored casing. + # Windows: case-insensitive lookup so C:\Models and c:\models dedup. if is_win: existing = conn.execute( "SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE", @@ -639,9 +910,8 @@ def add_scan_folder(path: str) -> dict: ) conn.commit() except sqlite3.IntegrityError: - pass # duplicate -- fall through to SELECT - # Use the same collation as the pre-check so we find the row even - # when a concurrent writer stored it with different casing (Windows). + pass # duplicate; fall through to SELECT + # Same collation as the pre-check to catch concurrent writes (Windows). fallback_sql = ( "SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE" if is_win @@ -681,6 +951,7 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict: "modelType": data["model_type"], "modelId": data.get("model_id") or "", "pairId": data.get("pair_id") or None, + "projectId": data.get("project_id") or None, "archived": bool(data["archived"]), "createdAt": data["created_at"], "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"), @@ -688,6 +959,21 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict: } +def _chat_project_from_row(row: sqlite3.Row) -> dict: + data = dict(row) + root_path = data.get("root_path") + return { + "id": data["id"], + "name": data["name"], + "instructions": data.get("instructions") or "", + "rootPath": root_path or None, + "sandboxPath": os.path.join(root_path, "sandbox") if root_path else None, + "archived": bool(data["archived"]), + "createdAt": data["created_at"], + "updatedAt": data["updated_at"], + } + + def _chat_message_from_row(row: sqlite3.Row) -> dict: data = dict(row) message = { @@ -713,13 +999,14 @@ def upsert_chat_thread(thread: dict) -> dict: conn.execute( """ INSERT INTO chat_threads - (id, title, model_type, model_id, pair_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, model_type = excluded.model_type, model_id = excluded.model_id, pair_id = excluded.pair_id, + project_id = excluded.project_id, archived = excluded.archived, created_at = excluded.created_at, openai_code_exec_container_id = excluded.openai_code_exec_container_id, @@ -731,6 +1018,7 @@ def upsert_chat_thread(thread: dict) -> dict: thread["modelType"], thread.get("modelId") or "", thread.get("pairId"), + thread.get("projectId"), 1 if thread.get("archived") else 0, int(thread["createdAt"]), thread.get("openaiCodeExecContainerId"), @@ -749,6 +1037,7 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]: "modelType": ("model_type", patch.get("modelType")), "modelId": ("model_id", patch.get("modelId")), "pairId": ("pair_id", patch.get("pairId")), + "projectId": ("project_id", patch.get("projectId")), "archived": ("archived", 1 if patch.get("archived") else 0), "createdAt": ("created_at", patch.get("createdAt")), "openaiCodeExecContainerId": ( @@ -794,6 +1083,7 @@ def get_chat_thread(id: str) -> Optional[dict]: def list_chat_threads( model_type: str | None = None, pair_id: str | None = None, + project_id: str | None = None, include_archived: bool = True, ) -> list[dict]: clauses = [] @@ -804,6 +1094,9 @@ def list_chat_threads( if pair_id is not None: clauses.append("pair_id = ?") values.append(pair_id) + if project_id is not None: + clauses.append("project_id = ?") + values.append(project_id) if not include_archived: clauses.append("archived = 0") where = f"WHERE {' AND '.join(clauses)}" if clauses else "" @@ -846,6 +1139,136 @@ def count_chat_threads() -> int: conn.close() +def upsert_chat_project(project: dict) -> dict: + existing = get_chat_project(project["id"]) + root_path = existing.get("rootPath") if existing else None + if not root_path: + root_path = _default_project_root(project) + root_path = _ensure_project_workspace(root_path) + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO chat_projects + (id, name, instructions, root_path, archived, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + name = excluded.name, + instructions = excluded.instructions, + root_path = COALESCE(chat_projects.root_path, excluded.root_path), + archived = excluded.archived, + created_at = excluded.created_at, + updated_at = excluded.updated_at + """, + ( + project["id"], + project["name"], + project.get("instructions") or "", + root_path, + 1 if project.get("archived") else 0, + int(project["createdAt"]), + int(project["updatedAt"]), + ), + ) + conn.commit() + return get_chat_project(project["id"]) or project + finally: + conn.close() + + +def update_chat_project(id: str, patch: dict) -> Optional[dict]: + allowed = { + "name": ("name", patch.get("name")), + "instructions": ("instructions", patch.get("instructions")), + "archived": ("archived", 1 if patch.get("archived") else 0), + "createdAt": ("created_at", patch.get("createdAt")), + "updatedAt": ("updated_at", patch.get("updatedAt")), + } + assignments = [] + values = [] + for key, (column, value) in allowed.items(): + if key in patch: + assignments.append(f"{column} = ?") + values.append(value) + if not assignments: + return get_chat_project(id) + + conn = get_connection() + try: + conn.execute( + f"UPDATE chat_projects SET {', '.join(assignments)} WHERE id = ?", + (*values, id), + ) + conn.commit() + row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone() + return _chat_project_from_row(row) if row is not None else None + finally: + conn.close() + + +def ensure_chat_project_workspace(id: str) -> Optional[dict]: + project = get_chat_project(id) + if project is None: + return None + root_path = project.get("rootPath") or _default_project_root(project) + root_path = _ensure_project_workspace(root_path) + if project.get("rootPath") == root_path: + return project + conn = get_connection() + try: + conn.execute( + "UPDATE chat_projects SET root_path = ? WHERE id = ?", + (root_path, id), + ) + conn.commit() + finally: + conn.close() + return get_chat_project(id) + + +def get_chat_project(id: str) -> Optional[dict]: + conn = get_connection() + try: + row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone() + return _chat_project_from_row(row) if row is not None else None + finally: + conn.close() + + +def list_chat_projects(include_archived: bool = False) -> list[dict]: + conn = get_connection() + try: + where = "" if include_archived else "WHERE archived = 0" + rows = conn.execute( + f"SELECT * FROM chat_projects {where} ORDER BY updated_at DESC" + ).fetchall() + return [_chat_project_from_row(row) for row in rows] + finally: + conn.close() + + +def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]: + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone() + if row is None: + conn.rollback() + return None + project = _chat_project_from_row(row) + conn.execute("DELETE FROM chat_threads WHERE project_id = ?", (id,)) + conn.execute("DELETE FROM chat_projects WHERE id = ?", (id,)) + conn.commit() + if delete_files: + _delete_project_workspace(project) + return project + except Exception: + conn.rollback() + raise + finally: + conn.close() + + class ChatMessageConflictError(RuntimeError): """Raised when a chat message id already belongs to another thread.""" @@ -866,9 +1289,7 @@ def _parse_chat_setting_json(key: str, value_json: str) -> tuple[bool, Any]: return False, None -def _load_chat_settings_for_merge( - conn: sqlite3.Connection, -) -> tuple[dict[str, Any], set[str]]: +def _load_chat_settings_for_merge(conn: sqlite3.Connection) -> tuple[dict[str, Any], set[str]]: rows = conn.execute("SELECT key, value_json FROM chat_settings").fetchall() current: dict[str, Any] = {} corrupt: set[str] = set() @@ -895,9 +1316,7 @@ def _load_chat_settings_for_merge( def _raise_if_chat_message_thread_conflicts( - conn: sqlite3.Connection, - thread_id: str, - message_ids: list[str], + conn: sqlite3.Connection, thread_id: str, message_ids: list[str] ) -> None: unique_ids = list(dict.fromkeys(message_ids)) if not unique_ids: @@ -1006,12 +1425,8 @@ def sync_chat_messages( m.get("parentId"), m["role"], json.dumps(m.get("content", [])), - json.dumps(m.get("attachments")) - if m.get("attachments") is not None - else None, - json.dumps(m.get("metadata")) - if m.get("metadata") is not None - else None, + json.dumps(m.get("attachments")) if m.get("attachments") is not None else None, + json.dumps(m.get("metadata")) if m.get("metadata") is not None else None, int(m["createdAt"]), ) for m in messages @@ -1088,12 +1503,44 @@ def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]: conn.close() +def get_app_setting(key: str, fallback = None): + conn = get_connection() + try: + row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone() + if row is None: + return fallback + return _json_loads(row["value_json"], fallback) + finally: + conn.close() + + +def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]: + if not settings: + return {} + conn = get_connection() + try: + now = datetime.now(timezone.utc).isoformat() + conn.executemany( + """ + INSERT INTO app_settings (key, value_json, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value_json = excluded.value_json, + updated_at = excluded.updated_at + """, + [(key, json.dumps(value), now) for key, value in settings.items()], + ) + conn.commit() + rows = conn.execute("SELECT key, value_json FROM app_settings ORDER BY key").fetchall() + return {row["key"]: _json_loads(row["value_json"], None) for row in rows} + finally: + conn.close() + + def list_chat_settings() -> dict[str, Any]: conn = get_connection() try: - rows = conn.execute( - "SELECT key, value_json FROM chat_settings ORDER BY key" - ).fetchall() + rows = conn.execute("SELECT key, value_json FROM chat_settings ORDER BY key").fetchall() settings: dict[str, Any] = {} for row in rows: settings[row["key"]] = _json_loads(row["value_json"], None) @@ -1124,9 +1571,7 @@ def upsert_chat_settings(settings: dict[str, Any]) -> dict[str, Any]: conn.close() -def _deep_merge_settings( - current: dict[str, Any], updates: dict[str, Any] -) -> dict[str, Any]: +def _deep_merge_settings(current: dict[str, Any], updates: dict[str, Any]) -> dict[str, Any]: merged = dict(current) for key, value in updates.items(): current_value = merged.get(key) @@ -1138,8 +1583,8 @@ def _deep_merge_settings( def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]: - """Atomic read-merge-write under BEGIN IMMEDIATE so two concurrent writers - cannot drop one another's updates.""" + """Atomic read-merge-write under BEGIN IMMEDIATE so concurrent writers + cannot drop each other's updates.""" if not updates: return list_chat_settings() conn = get_connection() @@ -1147,9 +1592,7 @@ def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]: conn.execute("BEGIN IMMEDIATE") current, corrupt = _load_chat_settings_for_merge(conn) unsafe_partial_keys = [ - key - for key, value in updates.items() - if key in corrupt and isinstance(value, dict) + key for key, value in updates.items() if key in corrupt and isinstance(value, dict) ] if unsafe_partial_keys: conn.commit() @@ -1187,16 +1630,10 @@ def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]: def list_chat_legacy_imports() -> list[str]: - """Return the legacy_thread_id of every thread already imported. - - Cheap: scans a single small PK-only table. The frontend stuffs the - result into a Set before walking Dexie, so the diff is O(|Dexie|). - """ + """Return the legacy_thread_id of every thread already imported.""" conn = get_connection() try: - rows = conn.execute( - "SELECT legacy_thread_id FROM chat_legacy_imports" - ).fetchall() + rows = conn.execute("SELECT legacy_thread_id FROM chat_legacy_imports").fetchall() return [row[0] for row in rows] finally: conn.close() @@ -1205,13 +1642,8 @@ def list_chat_legacy_imports() -> list[str]: def upsert_chat_legacy_imports(legacy_thread_ids: list[str]) -> tuple[int, int]: """Mark each given legacy thread id as imported. Idempotent. - Returns (accepted, inserted): - - accepted: number of non-empty deduped input ids - - inserted: number of rows that were actually new (not already in ledger) - - ON CONFLICT DO NOTHING keeps the existing imported_at when an id is - recorded twice. INSERT...RETURNING reports only the rows that were - actually inserted, so callers can distinguish first-time imports + Returns (accepted, inserted): count of deduped non-empty input ids, and + count of rows actually new. RETURNING lets callers tell first-time imports from idempotent re-runs without an extra SELECT. """ ids = list(dict.fromkeys(tid for tid in legacy_thread_ids if tid)) diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index 6aa6d314c1..b0b9ee309c 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -1,30 +1,14 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Shared pytest configuration for the backend test suite. +"""Shared pytest configuration for the backend test suite. -Responsibilities: -1. Put the backend root on sys.path so `from models.inference import ...` - (and similar flat imports) resolve in test modules — mirrors how the - app itself is launched. -2. Provide a hybrid ``studio_server`` session fixture for end-to-end tests - (see ``test_studio_api.py``). The fixture supports two invocation modes: - - a. **External server.** If ``UNSLOTH_E2E_BASE_URL`` is set, tests point - at an already-running Studio instance. ``UNSLOTH_E2E_API_KEY`` must - also be set. This is the fast-iteration mode: start the server once - with ``unsloth studio run ...``, then run pytest against it many - times with no per-run GGUF load cost. - - b. **Fixture-managed server.** Otherwise, the fixture launches a fresh - server via ``_start_server`` and tears it down at session end. This - is the one-shot mode for CI or a clean-slate verification run. - - The model / variant for mode (b) come from ``--unsloth-model`` / - ``--unsloth-gguf-variant`` pytest options, then ``UNSLOTH_E2E_MODEL`` / - ``UNSLOTH_E2E_VARIANT`` env vars, then the defaults in - ``test_studio_api.py``. +Puts the backend root on sys.path (mirrors app launch) and provides a hybrid +``studio_server`` session fixture for end-to-end tests with two modes: +external server (``UNSLOTH_E2E_BASE_URL``/``UNSLOTH_E2E_API_KEY``) for fast +iteration, or a fixture-managed server started/torn down per session for CI. +Model/variant for the managed mode resolve from ``--unsloth-model`` / +``--unsloth-gguf-variant``, then env vars, then ``test_studio_api.py`` defaults. """ import os @@ -33,13 +17,13 @@ from pathlib import Path import pytest -# Add backend root to sys.path (mirrors how the app itself is launched) +# Add backend root to sys.path (mirrors app launch) _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) -# ── Pytest CLI options ─────────────────────────────────────────────── +# Pytest CLI options def pytest_addoption(parser): @@ -71,25 +55,16 @@ def pytest_addoption(parser): ) -# ── E2E server fixtures ────────────────────────────────────────────── +# E2E server fixtures @pytest.fixture(scope = "session") def studio_server(request): """Yield ``(base_url, api_key)`` for e2e tests. - Resolution order: - - 1. If ``UNSLOTH_E2E_BASE_URL`` is set → point at that server, - require ``UNSLOTH_E2E_API_KEY`` alongside (skip if missing). - 2. Otherwise → start a fresh ``unsloth studio run`` subprocess via - the existing ``_start_server`` helper in ``test_studio_api.py`` - and tear it down on session teardown. - - Session-scoped so the expensive GGUF load happens at most once per - pytest invocation. Lazily instantiated — tests that don't request - the fixture (e.g. the unit tests in ``test_anthropic_messages.py`` - or ``test_help_output``) do not trigger server startup. + Uses ``UNSLOTH_E2E_BASE_URL`` (requires ``UNSLOTH_E2E_API_KEY``) if set, + else starts/tears down a fresh server via ``_start_server``. Session-scoped + and lazy so the GGUF load happens at most once and only when requested. """ external_url = os.environ.get("UNSLOTH_E2E_BASE_URL") if external_url: @@ -103,9 +78,7 @@ def studio_server(request): yield external_url, api_key return - # Lazy import: pytest has already loaded test_studio_api into - # sys.modules by the time any test requests this fixture, so this - # is a cache hit, not a re-execution. + # Lazy import; pytest has already loaded test_studio_api, so this is a cache hit. import test_studio_api as _e2e model = ( @@ -136,3 +109,71 @@ def base_url(studio_server): def api_key(studio_server): """API key for the e2e Studio server (from ``studio_server``).""" return studio_server[1] + + +# ── RAG fixtures ───────────────────────────────────────────────────── + + +@pytest.fixture +def rag_home(tmp_path, monkeypatch): + """Isolate the RAG database under a fresh UNSLOTH_STUDIO_HOME per test. + + Points the storage root at ``tmp_path`` and resets the lazy schema flag so + each test starts from an empty rag.db. Yields the temp home path. + """ + from storage import rag_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(rag_db, "_schema_ready", False) + return tmp_path + + +@pytest.fixture +def rag_conn(rag_home): + """A fresh RAG connection bound to the isolated ``rag_home`` database.""" + from storage import rag_db + + conn = rag_db.get_connection() + try: + yield conn + finally: + conn.close() + + +@pytest.fixture +def stub_embeddings(monkeypatch): + """Stub ``core.rag.embeddings`` with deterministic hash-based vectors. + + Lets store / retrieval / ingestion tests run fast without downloading a + sentence-transformers model. Returns the fixed embedding dimension. + """ + import hashlib + import math + + from core.rag import embeddings + + dim = 32 + + def _vec(text: str): + seed = hashlib.sha256(text.encode("utf-8")).digest() + raw = [seed[i % len(seed)] / 255.0 for i in range(dim)] + norm = math.sqrt(sum(x * x for x in raw)) or 1.0 + return [x / norm for x in raw] + + def fake_encode( + texts, + *, + model_name = None, + normalize = True, + ): + return [_vec(t) for t in texts] + + monkeypatch.setattr(embeddings, "encode", fake_encode) + monkeypatch.setattr(embeddings, "dim", lambda model_name = None: dim) + monkeypatch.setattr( + embeddings, + "token_counter", + lambda model_name = None: (lambda t: len(t.split())), + ) + monkeypatch.setattr(embeddings, "warm", lambda model_name = None: None) + return dim diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py new file mode 100644 index 0000000000..104182a232 --- /dev/null +++ b/studio/backend/tests/test_amd_apu_unified_memory.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs +(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS.""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.llama_cpp import LlamaCppBackend + + +def _fake_torch( + hip, + archs, + *, + cuda_ok = True, +): + t = types.ModuleType("torch") + t.version = types.SimpleNamespace(hip = hip) + t.cuda = types.SimpleNamespace( + is_available = lambda: cuda_ok, + device_count = lambda: len(archs), + get_device_properties = lambda i: types.SimpleNamespace(gcnArchName = archs[i]), + ) + return t + + +@pytest.mark.parametrize( + "hip,archs,expected", + [ + ("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped) + ("6.2.0", ["gfx1150"], True), # Strix Point APU + ("6.2.0", ["gfx1100"], False), # discrete RDNA3 + ("6.2.0", ["gfx1201"], False), # discrete RDNA4 + ("6.2.0", ["gfx942"], False), # MI300X (data center) + (None, ["sm_90"], False), # NVIDIA (no torch.version.hip) + ("6.2.0", ["gfx1100", "gfx1151"], True), # mixed dGPU + APU + ], +) +def test_apu_unified_memory_gating(monkeypatch, hip, archs, expected): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip, archs)) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is expected + + +def test_cpu_no_cuda_returns_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", [], cuda_ok = False)) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is False + + +def test_missing_torch_returns_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is False diff --git a/studio/backend/tests/test_anthropic_cache_ttl.py b/studio/backend/tests/test_anthropic_cache_ttl.py index e5d806e3c2..d39d0e9384 100644 --- a/studio/backend/tests/test_anthropic_cache_ttl.py +++ b/studio/backend/tests/test_anthropic_cache_ttl.py @@ -1,18 +1,12 @@ # 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 prompt_cache_ttl threading on the Anthropic path. +"""Unit tests for prompt_cache_ttl threading on the Anthropic path. -Anthropic accepts an optional ``ttl`` on each ``cache_control`` marker: -the default is the 5-minute ephemeral pool; ``ttl:"1h"`` writes into -the 1-hour pool instead. The 1h pool is the right pick when -conversations span multiple short bursts more than 5 minutes apart -- -1h writes are billed at 2x base input vs 1.25x for 5m, but reads stay -at 0.1x for both, so one extra read pays off the premium. - -These tests pin the outbound body shape: when prompt_cache_ttl="1h" -both cache_control markers carry ``ttl:"1h"``; default omits the field -entirely so the 5m pool is used; garbage values are silently dropped. +Anthropic's ``cache_control`` marker takes an optional ``ttl``: default 5m +pool, ``ttl:"1h"`` the 1h pool. These tests pin the outbound body shape: +"1h" puts ``ttl:"1h"`` on both markers; default omits the field; garbage +values are silently dropped. """ import asyncio @@ -128,13 +122,9 @@ def test_1h_ttl_writes_into_1h_pool(monkeypatch): def test_1h_ttl_does_not_send_extended_cache_ttl_beta_header(monkeypatch): - # The `extended-cache-ttl-2025-04-11` beta header that originally - # gated 1h cache TTL has been promoted to GA: verified live against - # api.anthropic.com on 2026-05-22 -- a request with - # `cache_control:{type:"ephemeral", ttl:"1h"}` and NO beta header - # returns 200 and populates `ephemeral_1h_input_tokens`. Pin the - # contract so we don't reintroduce the gate by accident; a future - # regression that re-adds the header would surface here. + # The extended-cache-ttl-2025-04-11 beta header is now GA (verified live + # 2026-05-22); 1h TTL works with no beta header. Pin so a regression that + # re-adds the header surfaces here. captured = _capture(monkeypatch, ttl = "1h") beta = captured["headers"].get("anthropic-beta", "") assert "extended-cache-ttl-2025-04-11" not in beta, beta @@ -155,8 +145,7 @@ def test_unknown_ttl_silently_dropped(monkeypatch, bogus): ccs = _cache_controls(captured["body"]) assert len(ccs) == 2, ccs for cc in ccs: - # Bogus TTLs must NOT round-trip; marker stays at the default - # (no `ttl` key, which means the 5m pool upstream). + # Bogus TTLs must not round-trip; marker stays at default (no ttl = 5m). assert cc == {"type": "ephemeral"}, cc diff --git a/studio/backend/tests/test_anthropic_citations.py b/studio/backend/tests/test_anthropic_citations.py index ab5ba10b56..25d5f5e886 100644 --- a/studio/backend/tests/test_anthropic_citations.py +++ b/studio/backend/tests/test_anthropic_citations.py @@ -3,11 +3,10 @@ """Tests for Anthropic ``citations_delta`` handling in the streaming proxy. -Verifies the proxy injects inline ``[N]`` markers after cited text, -dedupes by type-specific anchor (char_location, page_location, -content_block_location, search_result_location), forwards a synthetic -``document_citations`` tool_event at message_stop, and stays inert when -no citations_delta events fire. See +Verifies the proxy injects inline ``[N]`` markers after cited text, dedupes by +type-specific anchor (char_location, page_location, content_block_location, +search_result_location), forwards a synthetic ``document_citations`` tool_event +at message_stop, and stays inert when no citations_delta events fire. See https://platform.claude.com/docs/en/build-with-claude/citations """ @@ -127,8 +126,7 @@ def _joined(lines: list[str]) -> str: def test_no_citations_stream_unchanged(monkeypatch): - """Plain text streams pass through with no inline markers and no - document_citations tool_event.""" + """Plain text passes through with no inline markers or document_citations.""" lines = _capture( monkeypatch, [ diff --git a/studio/backend/tests/test_anthropic_citations_edge.py b/studio/backend/tests/test_anthropic_citations_edge.py index be1b5f7922..f89e5ebda8 100644 --- a/studio/backend/tests/test_anthropic_citations_edge.py +++ b/studio/backend/tests/test_anthropic_citations_edge.py @@ -50,10 +50,9 @@ def _capture( messages: list[dict] | None = None, captured_body: dict | None = None, ) -> list[str]: - """Drive ``stream_chat_completion`` against a mocked Anthropic - response and return the SSE lines. Pass ``captured_body`` to also - capture the outgoing request body for assertions on the translated - Anthropic shape. + """Drive ``stream_chat_completion`` against a mocked Anthropic response + and return the SSE lines. Pass ``captured_body`` to also capture the + outgoing request body for assertions on the translated Anthropic shape. """ def handler(request: httpx.Request) -> httpx.Response: @@ -80,8 +79,7 @@ def _capture( client = _make_client() try: async for line in client.stream_chat_completion( - messages = messages - or [{"role": "user", "content": "what color is grass?"}], + messages = messages or [{"role": "user", "content": "what color is grass?"}], model = "claude-opus-4-7", max_tokens = 64, ): @@ -148,8 +146,8 @@ def _joined(lines: list[str]) -> str: def _citation_payload(body: str) -> dict: - """Pull the ``document_citations`` synthetic tool_event from the - SSE body and return its payload. Raises if absent.""" + """Return the ``document_citations`` synthetic tool_event payload from + the SSE body. Raises if absent.""" assert "document_citations" in body, body for line in body.splitlines(): if not line.startswith("data: "): @@ -159,10 +157,7 @@ def _citation_payload(body: str) -> dict: except json.JSONDecodeError: continue tool_event = payload.get("_toolEvent") if isinstance(payload, dict) else None - if ( - isinstance(tool_event, dict) - and tool_event.get("type") == "document_citations" - ): + if isinstance(tool_event, dict) and tool_event.get("type") == "document_citations": return tool_event raise AssertionError("document_citations event not parsed out of SSE body") @@ -171,8 +166,8 @@ def _citation_payload(body: str) -> dict: def test_citation_with_no_preceding_text_still_emits_marker(monkeypatch): - """citations_delta before any text_delta must not crash; marker - lands at the start of the block.""" + """citations_delta before any text_delta must not crash; marker lands + at the start of the block.""" cit = { "type": "char_location", "document_index": 0, @@ -199,8 +194,8 @@ def test_citation_with_no_preceding_text_still_emits_marker(monkeypatch): def test_citations_delta_with_non_dict_citation_is_ignored(monkeypatch): - """Non-dict ``delta.citation`` must not crash, emit a marker, or - poison the document_citations list.""" + """Non-dict ``delta.citation`` must not crash, emit a marker, or poison + the document_citations list.""" lines = _capture( monkeypatch, [ @@ -224,8 +219,8 @@ def test_citations_delta_with_non_dict_citation_is_ignored(monkeypatch): def test_citations_delta_with_missing_citation_field_is_ignored(monkeypatch): - """Missing ``citation`` field is treated like a non-dict citation: - skip without crashing.""" + """Missing ``citation`` field is treated like a non-dict citation: skip + without crashing.""" lines = _capture( monkeypatch, [ @@ -249,8 +244,8 @@ def test_citations_delta_with_missing_citation_field_is_ignored(monkeypatch): def test_char_location_with_reversed_indices_does_not_crash(monkeypatch): - """Malformed char_location with reversed indices must not crash; - the dedup key accepts any int pair and still surfaces a footnote.""" + """Malformed char_location with reversed indices must not crash; the + dedup key accepts any int pair and still surfaces a footnote.""" cit = { "type": "char_location", "document_index": 0, @@ -279,8 +274,8 @@ def test_char_location_with_reversed_indices_does_not_crash(monkeypatch): def test_page_location_missing_document_index_does_not_crash(monkeypatch): - """page_location missing ``document_index`` still produces a - footnote; dedup key falls back to ``None`` for the missing field.""" + """page_location missing ``document_index`` still produces a footnote; + dedup key falls back to ``None`` for the missing field.""" cit = { "type": "page_location", "document_title": "Untitled PDF", @@ -307,8 +302,8 @@ def test_page_location_missing_document_index_does_not_crash(monkeypatch): def test_content_block_location_with_non_int_block_index_does_not_crash(monkeypatch): - """content_block_location with string block indices must not crash; - dedup key tolerates non-int values.""" + """content_block_location with string block indices must not crash; dedup + key tolerates non-int values.""" cit = { "type": "content_block_location", "document_index": 0, @@ -336,8 +331,8 @@ def test_content_block_location_with_non_int_block_index_does_not_crash(monkeypa def test_unknown_citation_type_falls_back_to_stringified_key(monkeypatch): - """Unknown citation ``type`` (forward-compat) still dedupes: - identical ones collapse, differing ones get distinct numbers.""" + """Unknown citation ``type`` (forward-compat) still dedupes: identical + ones collapse, differing ones get distinct numbers.""" cit_a = { "type": "future_shape_location", "anchor": "abc", @@ -373,8 +368,8 @@ def test_unknown_citation_type_falls_back_to_stringified_key(monkeypatch): def test_mixed_citation_types_same_document_get_distinct_keys(monkeypatch): - """char_location and page_location on the same document_index are - distinct shapes; dedup key uses citation type as its first slot.""" + """char_location and page_location on the same document_index are distinct + shapes; dedup key uses citation type as its first slot.""" cit_char = { "type": "char_location", "document_index": 0, @@ -410,9 +405,9 @@ def test_mixed_citation_types_same_document_get_distinct_keys(monkeypatch): def test_cited_text_is_preserved_in_synthetic_event(monkeypatch): - """``cited_text`` must survive into the synthetic event so the - Sources panel can render it as a tooltip. Anthropic does not bill - cited_text against output tokens, so preserving it is free.""" + """``cited_text`` must survive into the synthetic event so the Sources + panel can render it as a tooltip. Anthropic does not bill cited_text + against output tokens, so preserving it is free.""" cit = { "type": "char_location", "document_index": 0, @@ -439,8 +434,8 @@ def test_cited_text_is_preserved_in_synthetic_event(monkeypatch): def test_internal_key_field_never_leaks_to_client(monkeypatch): - """The internal ``_key`` dedup sentinel must be stripped before - the synthetic event is forwarded; it is not an Anthropic field.""" + """The internal ``_key`` dedup sentinel must be stripped before the + synthetic event is forwarded; it is not an Anthropic field.""" cit = { "type": "char_location", "document_index": 0, @@ -469,8 +464,8 @@ def test_internal_key_field_never_leaks_to_client(monkeypatch): def test_citation_across_multiple_content_blocks_numbers_continue(monkeypatch): - """Footnote numbering is per-message, not per-content-block: - citations across separate blocks emit [1] then [2].""" + """Footnote numbering is per-message, not per-content-block: citations + across separate blocks emit [1] then [2].""" cit_a = { "type": "char_location", "document_index": 0, @@ -513,9 +508,9 @@ def test_citation_across_multiple_content_blocks_numbers_continue(monkeypatch): def test_inline_marker_lands_after_text_run(monkeypatch): - """Inline ``[N]`` must land AFTER the cited text run: Anthropic - streams text then citation, so the proxy emits ``"...green.[1]"`` - not ``"[1]green"``.""" + """Inline ``[N]`` must land AFTER the cited text run: Anthropic streams + text then citation, so the proxy emits ``"...green.[1]"`` not + ``"[1]green"``.""" cit = { "type": "char_location", "document_index": 0, @@ -545,8 +540,8 @@ def test_inline_marker_lands_after_text_run(monkeypatch): def test_no_synthetic_event_when_only_text_deltas(monkeypatch): - """No citations_delta means no synthetic ``document_citations`` - event; Sources panel relies on absence to suppress the section.""" + """No citations_delta means no synthetic ``document_citations`` event; + Sources panel relies on absence to suppress the section.""" lines = _capture( monkeypatch, [ @@ -565,9 +560,9 @@ def test_no_synthetic_event_when_only_text_deltas(monkeypatch): def test_input_document_translation_enables_citations(monkeypatch): - """``input_document`` must translate to an Anthropic ``document`` - block carrying ``citations: {enabled: true}`` (both base64 and url - source branches) so upstream emits citations_delta.""" + """``input_document`` must translate to an Anthropic ``document`` block + carrying ``citations: {enabled: true}`` (both base64 and url source + branches) so upstream emits citations_delta.""" captured_b64: dict = {} _capture( monkeypatch, @@ -635,8 +630,8 @@ def test_input_document_translation_enables_citations(monkeypatch): def test_cited_text_truncated_in_synthetic_event(monkeypatch): - """``cited_text`` is capped server-side so multi-KB spans do not - balloon the SSE payload.""" + """``cited_text`` is capped server-side so multi-KB spans don't balloon + the SSE payload.""" from core.inference.external_provider import _CITED_TEXT_MAX_LEN long_quote = "x" * (_CITED_TEXT_MAX_LEN + 4000) diff --git a/studio/backend/tests/test_anthropic_code_execution.py b/studio/backend/tests/test_anthropic_code_execution.py index 7f6fe58329..22bbb19125 100644 --- a/studio/backend/tests/test_anthropic_code_execution.py +++ b/studio/backend/tests/test_anthropic_code_execution.py @@ -1,29 +1,9 @@ # 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``. -""" +"""Tests for Anthropic server-side `code_execution` tool translation in +`_stream_anthropic`: request body/beta header, combined web_search, +bash/text_editor SSE tool events, and the container_expired error path.""" import asyncio import json @@ -116,16 +96,11 @@ def test_code_execution_tool_appended_to_request_body(monkeypatch): body = captured["body"] tools = body.get("tools") or [] - # Opus 4.7 gets the newer date-pinned variant (`_20260120`) that - # supports REPL state persistence + programmatic tool calling. - assert { - "type": "code_execution_20260120", - "name": "code_execution", - } in tools + # Opus 4.7 gets the newer date-pinned variant with REPL persistence. + assert {"type": "code_execution_20260120", "name": "code_execution"} in tools # No web_search entry when only code_execution is enabled. assert all("web_search" not in (t.get("type") or "") for t in tools) - # Beta header still carries the documented flag; both `_20250825` - # and `_20260120` are unlocked by the same header per upstream docs. + # Same beta header unlocks both _20250825 and _20260120. beta_header = captured["headers"].get("anthropic-beta", "") assert "code-execution-2025-08-25" in beta_header @@ -196,14 +171,10 @@ def test_no_code_execution_tool_when_pill_off(monkeypatch): _drive(run()) tools = captured["body"].get("tools") or [] - # Pill off -- neither the legacy nor the new code_execution variant - # may appear on the wire. + # Pill off: no code_execution variant on the wire. assert all("code_execution" not in (t.get("type") or "") 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", "" - ) + # Beta header must omit code-execution when the tool is off (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): @@ -275,7 +246,8 @@ def test_bash_code_execution_emits_tool_start_and_end(monkeypatch): 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"} + # `_server_tool: True` marks a provider-side synthetic tool card. + assert start["arguments"] == {"kind": "bash", "command": "ls -la", "_server_tool": True} assert end["type"] == "tool_end" assert end["tool_call_id"] == "srvtoolu_1" @@ -302,8 +274,7 @@ def test_text_editor_create_emits_kind_and_status(monkeypatch): "delta": { "type": "input_json_delta", "partial_json": ( - '{"command": "create", "path": "new_file.txt", ' - '"file_text": "hi"}' + '{"command": "create", "path": "new_file.txt", "file_text": "hi"}' ), }, }, diff --git a/studio/backend/tests/test_anthropic_compaction.py b/studio/backend/tests/test_anthropic_compaction.py index 92b9280146..1528eebe8b 100644 --- a/studio/backend/tests/test_anthropic_compaction.py +++ b/studio/backend/tests/test_anthropic_compaction.py @@ -3,16 +3,13 @@ """Unit tests for Anthropic server-side context compaction wiring. -Compaction is a beta feature (header ``compact-2026-01-12``) gated to -Opus 4.6, Opus 4.7, Sonnet 4.6, and Mythos preview. When enabled, -Studio attaches ``context_management.edits[{type:"compact_20260112", -trigger:{type:"input_tokens", value:N}}]`` to the outbound body. The -minimum upstream-accepted threshold is 50k tokens; lower values are -clamped to 50k so the request doesn't 400. +Compaction is a beta (header ``compact-2026-01-12``) gated to Opus 4.6/4.7, +Sonnet 4.6, and Mythos preview. When enabled, Studio attaches +``context_management.edits[{type:"compact_20260112", trigger:{type:"input_tokens", +value:N}}]``; the 50k-token minimum is clamped up so the request doesn't 400. -These tests pin: the body shape per model, the beta header merge with -the existing code-execution beta, threshold clamping, and silent no-op -on unsupported models. +Pins body shape per model, beta header merge with code-execution, threshold +clamping, and silent no-op on unsupported models. """ import asyncio @@ -40,7 +37,12 @@ def _make_client() -> ExternalProviderClient: ) -def _capture(monkeypatch, model: str, threshold, tools = None) -> dict: +def _capture( + monkeypatch, + model: str, + threshold, + tools = None, +) -> dict: captured: dict = {} def handler(request: httpx.Request) -> httpx.Response: @@ -120,13 +122,9 @@ def test_supported_model_attaches_compaction_block_and_beta(monkeypatch): def test_threshold_clamped_to_50k_minimum(monkeypatch): # Below-min values get clamped UP so we don't 400 upstream. captured = _capture(monkeypatch, "claude-opus-4-7", 60_000) - assert ( - captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 60_000 - ) + assert captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 60_000 captured = _capture(monkeypatch, "claude-opus-4-7", 1) - assert ( - captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 50_000 - ) + assert captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 50_000 # ── beta header merge with code execution ──────────────────────────── @@ -151,10 +149,7 @@ def test_unsupported_model_silently_drops_compaction(monkeypatch): captured = _capture(monkeypatch, "claude-haiku-4-5-20251001", 150_000) assert "context_management" not in captured["body"] # The beta header must not carry compact-2026-01-12 either. - assert "compact-2026-01-12" not in captured["headers"].get( - "anthropic-beta", - "", - ) + assert "compact-2026-01-12" not in captured["headers"].get("anthropic-beta", "") # ── omitted threshold leaves body untouched ───────────────────────── @@ -163,20 +158,15 @@ def test_unsupported_model_silently_drops_compaction(monkeypatch): def test_omitted_threshold_no_body_field(monkeypatch): captured = _capture(monkeypatch, "claude-opus-4-7", None) assert "context_management" not in captured["body"] - assert "compact-2026-01-12" not in captured["headers"].get( - "anthropic-beta", - "", - ) + assert "compact-2026-01-12" not in captured["headers"].get("anthropic-beta", "") # ── ChatCompletionRequest schema accepts sub-50k threshold ────────── def test_chat_completion_request_accepts_sub_50k_compaction_threshold(): - # Codex P1 caught that ge=50_000 on the field caused FastAPI to - # 422 the request before the in-helper clamp could fire. The - # schema must accept any positive int and let _stream_anthropic - # clamp upward. + # ge=50_000 on the field would 422 before the in-helper clamp fires; the + # schema must accept any positive int and let _stream_anthropic clamp up. from models.inference import ChatCompletionRequest req = ChatCompletionRequest.model_validate( @@ -212,17 +202,11 @@ def test_chat_completion_request_accepts_sub_50k_compaction_threshold(): # ── usage.iterations[] surfaces compaction tokens ────────────────── -def test_message_delta_iterations_array_aggregates_compaction_tokens( - monkeypatch, capsys -): - # When Anthropic compacts mid-stream, the SSE message_delta usage - # payload carries `iterations: [{type:"compaction", ...}, ...]`. - # The top-level input_tokens / output_tokens only account for the - # `message` iteration, so the cost surface needs the compaction - # totals exposed separately. The stream helper folds them into - # last_usage as `compaction_input_tokens` / `compaction_output_tokens` - # and surfaces them in the closing summary log so an operator can - # eyeball "did compaction cost us 180k tokens this turn?". +def test_message_delta_iterations_array_aggregates_compaction_tokens(monkeypatch, capsys): + # On mid-stream compaction the message_delta usage carries + # `iterations: [{type:"compaction", ...}, ...]`. Top-level tokens only cover + # the `message` iteration, so the helper folds compaction totals into + # last_usage and surfaces them in the closing summary log. def http_handler(request: httpx.Request) -> httpx.Response: body = ( @@ -265,8 +249,7 @@ def test_message_delta_iterations_array_aggregates_compaction_tokens( _drive(run()) - # structlog renders the closing summary through the stdlib bridge, - # which lands on stdout. Capture and check the rendered line. + # structlog renders the closing summary onto stdout; check the rendered line. out = capsys.readouterr().out summary = next( (line for line in out.splitlines() if "Anthropic stream complete" in line), @@ -277,9 +260,8 @@ def test_message_delta_iterations_array_aggregates_compaction_tokens( def test_message_delta_no_iterations_leaves_compaction_keys_unset(monkeypatch, capsys): - # Re-applying a previous compaction block does NOT emit a fresh - # iterations array. The helper must not invent compaction keys - # in that case (would otherwise double-bill). + # Re-applying a prior compaction block emits no fresh iterations array; + # the helper must not invent compaction keys (would double-bill). def http_handler(request: httpx.Request) -> httpx.Response: body = ( b"event: message_delta\n" @@ -338,19 +320,14 @@ def _async_collect(agen): def test_compaction_block_emitted_as_tool_event(monkeypatch): - # Codex P1: once context_management is enabled and Anthropic runs - # compaction during a turn, the response carries a - # `{type:"compaction", content:""}` block. The translator - # must surface it so the chat-adapter can persist it onto the - # assistant message; otherwise the next turn loses the state and - # Anthropic re-compacts from scratch. + # When Anthropic compacts during a turn, the response carries a + # `{type:"compaction", content:""}` block. The translator must surface + # it so the chat-adapter persists it; else the next turn loses state and re-compacts. def http_handler(request: httpx.Request) -> httpx.Response: - # Anthropic ships compaction blocks as a content_block_start - # with `type:"compaction"`, then either includes the summary - # on that start event AND/OR streams it via text_delta events - # on the same block index. Test the streamed-delta path since - # it's the harder case. + # Compaction blocks arrive as a content_block_start with type:"compaction", + # with the summary on the start event AND/OR streamed via text_delta on the + # same index. Test the streamed-delta path (harder case). body = ( b"event: message_start\n" b'data: {"type":"message_start","message":{"usage":{}}}\n\n' @@ -417,9 +394,7 @@ def test_compaction_block_emitted_as_tool_event(monkeypatch): parsed = json.loads(raw) except json.JSONDecodeError: continue - # tool_event payloads ride inside chat.completion.chunk.choices[0].delta.content - # as a JSON-encoded string. The simpler path: look for the - # marker substring anywhere in the chunk. + # tool_event payloads ride inside the chunk as JSON; just match the marker. if "compaction_block" in raw: events.append(raw) assert events, f"no compaction_block tool event found in {lines}" @@ -452,10 +427,8 @@ def test_compaction_block_emitted_as_tool_event(monkeypatch): def test_compaction_block_round_trips_through_outbound_messages(monkeypatch): - # Once the prior turn persisted a compaction block onto the - # assistant message, the next turn's outbound body must forward - # the {type:"compaction", content:"..."} block to Anthropic - # verbatim so the API recognises the existing state. + # The next turn's outbound body must forward a persisted + # {type:"compaction", content:"..."} block verbatim so the API recognises the state. captured: dict = {} def http_handler(request: httpx.Request) -> httpx.Response: @@ -513,9 +486,7 @@ def test_compaction_block_round_trips_through_outbound_messages(monkeypatch): def test_compaction_content_part_accepted_by_chat_message_schema(): - # Without this Pydantic Tag the discriminated Union would 422 the - # request at parse time and the round-trip would never reach the - # translator. + # Without the Pydantic Tag the discriminated Union would 422 at parse time. from models.inference import ChatMessage msg = ChatMessage.model_validate( @@ -534,10 +505,8 @@ def test_compaction_content_part_accepted_by_chat_message_schema(): def test_build_external_messages_passes_compaction_for_anthropic_only(): - # Compaction is an Anthropic-only synthetic content part. The - # builder MUST gate it on provider_type=="anthropic"; every other - # provider would 400 on the unknown content type via generic - # /chat/completions passthrough (Codex P1 follow-up). + # Compaction is Anthropic-only; the builder must gate it on + # provider_type=="anthropic" since others 400 on the unknown content type. from models.inference import ChatMessage from routes.inference import _build_external_messages @@ -552,9 +521,7 @@ def test_build_external_messages_passes_compaction_for_anthropic_only(): } ) ] - out = _build_external_messages( - msgs, supports_vision = True, provider_type = "anthropic" - ) + out = _build_external_messages(msgs, supports_vision = True, provider_type = "anthropic") assert len(out) == 1 parts = out[0]["content"] assert parts[0] == {"type": "compaction", "content": "prior summary"} @@ -562,11 +529,9 @@ def test_build_external_messages_passes_compaction_for_anthropic_only(): def test_build_external_messages_strips_compaction_for_non_anthropic_providers(): - # Provider switch (or reused history) hands compaction blocks to a - # non-Anthropic provider. Those land on generic /chat/completions - # passthrough where the unknown content type fails the upstream - # validator. Builder must strip the part for every non-anthropic - # provider, including OpenAI/DeepSeek/Mistral/Gemini/Kimi/OpenRouter. + # A provider switch or reused history can hand compaction blocks to a + # non-Anthropic provider whose validator rejects the unknown type, so the + # builder must strip the part for every non-anthropic provider. from models.inference import ChatMessage from routes.inference import _build_external_messages @@ -582,9 +547,7 @@ def test_build_external_messages_strips_compaction_for_non_anthropic_providers() ) ] for provider in ("openai", "deepseek", "mistral", "gemini", "kimi", "openrouter"): - out = _build_external_messages( - msgs, supports_vision = True, provider_type = provider - ) + out = _build_external_messages(msgs, supports_vision = True, provider_type = provider) assert len(out) == 1, (provider, out) parts = out[0]["content"] types = [p.get("type") for p in parts if isinstance(p, dict)] @@ -594,9 +557,7 @@ def test_build_external_messages_strips_compaction_for_non_anthropic_providers() def test_build_external_messages_strips_compaction_when_provider_type_unknown(): - # Defensive: if provider_type is None (legacy path) the part must - # also be stripped -- forwarding to an unknown destination is - # never safe. + # Defensive: provider_type=None (legacy path) must also strip the part. from models.inference import ChatMessage from routes.inference import _build_external_messages @@ -618,9 +579,8 @@ def test_build_external_messages_strips_compaction_when_provider_type_unknown(): def test_build_external_messages_non_vision_anthropic_keeps_compaction(): - # Defensive: even though compaction-capable Anthropic models all - # currently report supports_vision=True, gate the non-vision branch - # by provider_type too so future config changes don't drop it. + # Defensive: gate the non-vision branch by provider_type too, so future + # config changes don't drop compaction for Anthropic. from models.inference import ChatMessage from routes.inference import _build_external_messages @@ -635,14 +595,10 @@ def test_build_external_messages_non_vision_anthropic_keeps_compaction(): } ) ] - out = _build_external_messages( - msgs, supports_vision = False, provider_type = "anthropic" - ) + out = _build_external_messages(msgs, supports_vision = False, provider_type = "anthropic") parts = out[0]["content"] assert {"type": "compaction", "content": "prior summary"} in parts # Non-anthropic + non-vision -> compaction stripped, text collapsed # back to a string. - out2 = _build_external_messages( - msgs, supports_vision = False, provider_type = "deepseek" - ) + out2 = _build_external_messages(msgs, supports_vision = False, provider_type = "deepseek") assert out2[0]["content"] == "answer", out2 diff --git a/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py b/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py index e7e5ec64d4..da0c0941ca 100644 --- a/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py +++ b/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py @@ -3,9 +3,9 @@ """Tests for Anthropic fast-mode wiring and streaming refusal handling. -fast_mode=True on Opus 4.6/4.7 attaches the ``fast-mode-2026-02-01`` -beta header and sets ``speed: "fast"``; unsupported models drop both. -Streaming ``stop_reason: "refusal"`` surfaces a user notice before the +fast_mode=True on Opus 4.6/4.7 attaches the ``fast-mode-2026-02-01`` beta +header and sets ``speed: "fast"``; unsupported models drop both. Streaming +``stop_reason: "refusal"`` surfaces a user notice before the ``content_filter`` finish chunk. https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals """ @@ -58,8 +58,12 @@ def _refusal_sse() -> bytes: ) -def _capture(monkeypatch, sse: bytes = b"", **kwargs) -> tuple[dict, list[str]]: - """Install a MockTransport, drive one streamed call, return body+lines.""" +def _capture( + monkeypatch, + sse: bytes = b"", + **kwargs, +) -> tuple[dict, list[str]]: + """Install a MockTransport, drive one streamed call; return body+lines.""" captured: dict = {} def handler(request: httpx.Request) -> httpx.Response: @@ -153,12 +157,12 @@ def test_refusal_emits_user_facing_notice_and_content_filter_finish(monkeypatch) def test_refusal_emits_tool_event_for_chat_adapter_drop(monkeypatch): """Refused turns emit an out-of-band `_toolEvent` that the chat-adapter - latches into assistant `metadata.custom.anthropicRefusal`, driving - the next-request prune. Tool event (not text) prevents spoofing. + latches into assistant `metadata.custom.anthropicRefusal`, driving the + next-request prune. Tool event (not text) prevents spoofing. """ _, lines = _capture(monkeypatch, sse = _refusal_sse()) body = "\n".join(lines) assert '"_toolEvent": {"type": "anthropic_refusal"}' in body, body - # Visible refusal text must not embed a sentinel that could spoof - # a context reset if echoed by another assistant message. + # Visible refusal text must not embed a sentinel that could spoof a + # context reset if echoed by another assistant message. assert "studio:anthropic-refusal" not in body, body diff --git a/studio/backend/tests/test_anthropic_fast_mode_edge.py b/studio/backend/tests/test_anthropic_fast_mode_edge.py index 0052cb94ad..dd69d77590 100644 --- a/studio/backend/tests/test_anthropic_fast_mode_edge.py +++ b/studio/backend/tests/test_anthropic_fast_mode_edge.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Edge-case coverage for the Anthropic fast-mode + refusal wiring. +"""Edge-case coverage for Anthropic fast-mode + refusal wiring. -Complements ``test_anthropic_fast_mode_and_refusal.py`` (happy path) -with dated snapshots, strict opt-in (future Opus families do not -auto-enable), multi-beta header merging, refusal stream ordering, and -the non-destruction guarantee for unset/None fast_mode. +Complements ``test_anthropic_fast_mode_and_refusal.py`` (happy path) with +dated snapshots, strict opt-in (future Opus families do not auto-enable), +multi-beta header merging, refusal stream ordering, and the +non-destruction guarantee for unset/None fast_mode. """ import asyncio @@ -60,7 +60,11 @@ def _refusal_sse(model: str = "claude-opus-4-7") -> bytes: ) -def _capture(monkeypatch, sse: bytes = b"", **kwargs) -> tuple[dict, list[str]]: +def _capture( + monkeypatch, + sse: bytes = b"", + **kwargs, +) -> tuple[dict, list[str]]: """Install a MockTransport, drive one streamed call, return body+lines.""" captured: dict = {} @@ -124,7 +128,7 @@ def test_fast_mode_attaches_on_dated_opus_4_6_snapshot(monkeypatch): # ──────────────────────────── strict opt-in semantics ──────────────────────────── def test_fast_mode_does_not_auto_enable_on_future_opus_4_8(monkeypatch): - """Future ``claude-opus-4-8`` must not auto-enable; opt-in per family.""" + """Future ``claude-opus-4-8`` must not auto-enable; per-family opt-in.""" cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-8") assert "speed" not in cap["body"], cap["body"] assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") @@ -244,7 +248,7 @@ def test_fast_mode_unset_is_byte_identical_to_omitted(monkeypatch): _drive(run()) assert cap_none["body"] == captured["body"], (cap_none["body"], captured["body"]) - # Headers can vary by httpx-injected fields (host, connection); compare + # Headers vary by httpx-injected fields (host, connection); compare # the load-bearing ones. for key in ("anthropic-version", "x-api-key", "content-type"): assert cap_none["headers"].get(key) == captured["headers"].get(key), key @@ -266,9 +270,7 @@ def test_refusal_notice_appears_before_content_filter_chunk(monkeypatch): """The notice content delta must precede the finish_reason chunk.""" _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") notice_idx = next(i for i, l in enumerate(lines) if "stopped by Anthropic" in l) - filter_idx = next( - i for i, l in enumerate(lines) if '"finish_reason": "content_filter"' in l - ) + filter_idx = next(i for i, l in enumerate(lines) if '"finish_reason": "content_filter"' in l) assert notice_idx < filter_idx, (notice_idx, filter_idx, lines) @@ -321,8 +323,7 @@ def test_refusal_chunk_is_proper_openai_delta_shape(monkeypatch): assert notice_chunk is not None, lines choice = notice_chunk["choices"][0] assert "delta" in choice and "content" in choice["delta"], notice_chunk - # Must NOT carry a finish_reason itself -- that comes on the next - # chunk. + # Must NOT carry a finish_reason itself -- that comes on the next chunk. assert choice.get("finish_reason") in (None,), notice_chunk # Refusal text is plain-spoken; no embedded sentinel. assert "studio:anthropic-refusal" not in choice["delta"]["content"] @@ -350,7 +351,6 @@ def test_fast_mode_prefix_tuple_matches_capability_doc(monkeypatch): """Tuple must exactly match the two families in the upstream docs: https://platform.claude.com/docs/en/build-with-claude/fast-mode.""" from core.inference.external_provider import _ANTHROPIC_FAST_MODE_PREFIXES - assert set(_ANTHROPIC_FAST_MODE_PREFIXES) == { "claude-opus-4-7", "claude-opus-4-6", @@ -421,9 +421,7 @@ def test_usage_speed_propagates_to_final_usage_chunk_fast(monkeypatch): def test_usage_speed_propagates_to_final_usage_chunk_standard(monkeypatch): _, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "standard")) parsed = [ - json.loads(l[len("data: ") :]) - for l in lines - if l.startswith("data: ") and '"usage"' in l + json.loads(l[len("data: ") :]) for l in lines if l.startswith("data: ") and '"usage"' in l ] speeds = [p["usage"].get("speed") for p in parsed if "usage" in p] assert "standard" in speeds, parsed @@ -433,9 +431,7 @@ def test_usage_speed_absent_when_anthropic_does_not_report(monkeypatch): """Studio must not invent ``usage.speed`` when upstream omits it.""" _, lines = _capture(monkeypatch) parsed = [ - json.loads(l[len("data: ") :]) - for l in lines - if l.startswith("data: ") and '"usage"' in l + json.loads(l[len("data: ") :]) for l in lines if l.startswith("data: ") and '"usage"' in l ] for p in parsed: usage = p.get("usage") or {} diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 842429d5af..e1230ae113 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1,15 +1,14 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -""" -Tests for the Anthropic Messages API schemas and translation layer. -No running server or GPU required. -""" +"""Tests for Anthropic Messages API schemas and translation layer (no server/GPU).""" import sys import os import json +import threading +import httpx import pytest _backend = os.path.join(os.path.dirname(__file__), "..") @@ -35,9 +34,12 @@ from core.inference.anthropic_compat import ( AnthropicPassthroughEmitter, ) from routes.inference import ( + _build_tool_action_nudge, _normalize_anthropic_openai_images, _select_anthropic_server_tools, _anthropic_requested_studio_tools, + _anthropic_passthrough_stream, + _anthropic_tool_non_streaming, anthropic_messages, ) from state.tool_policy import reset_tool_policy, set_tool_policy @@ -48,6 +50,51 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +# ===================================================================== +# Tool nudge tests +# ===================================================================== + + +class TestToolActionNudge: + def test_balanced_nudge_uses_expanded_web_and_code_tips(self): + nudge = _build_tool_action_nudge( + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + model_name = "Llama-3.1-70B-Instruct", + ) + + assert nudge.startswith("The current date is ") + assert "Tools are available when they materially improve" in nudge + assert "prefer using tools rather than answering from memory" not in nudge + assert "fetch its full content by calling web_search with the url parameter" in nudge + assert "Use code execution for math" in nudge + assert "render_html" not in nudge + + def test_balanced_nudge_preserves_compact_web_tip_and_artifact_gate(self): + nudge = _build_tool_action_nudge( + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "render_html"}}, + ], + model_name = "Llama-3.1-8B-Instruct", + ) + + assert "When using web_search, do not repeat the same search query." in nudge + assert "fetch its full content" not in nudge + assert "call render_html once" in nudge + + def test_balanced_nudge_empty_without_known_tool_categories(self): + assert ( + _build_tool_action_nudge( + tools = [], + model_name = "Llama-3.1-8B-Instruct", + ) + == "" + ) + + # ===================================================================== # Pydantic model tests # ===================================================================== @@ -77,6 +124,51 @@ class TestAnthropicModels: ) assert req.system == "You are helpful." + def test_system_role_message_normalized_to_system_field(self): + req = AnthropicMessagesRequest( + max_tokens = 50, + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hi"}, + ], + ) + assert req.system == "You are helpful." + assert len(req.messages) == 1 + assert req.messages[0].role == "user" + + def test_system_role_message_merges_with_existing_system_field(self): + req = AnthropicMessagesRequest( + max_tokens = 50, + system = "Base instructions.", + messages = [ + {"role": "user", "content": "Hi"}, + {"role": "system", "content": "Additional instructions."}, + {"role": "assistant", "content": "Hello."}, + ], + ) + assert req.system == "Base instructions.\n\nAdditional instructions." + assert [msg.role for msg in req.messages] == ["user", "assistant"] + + def test_system_role_message_with_null_content_ignored(self): + req = AnthropicMessagesRequest( + max_tokens = 50, + system = "Base.", + messages = [ + {"role": "system", "content": None}, + { + "role": "system", + "content": [ + None, + {"type": "text", "text": "Use short answers."}, + ], + }, + {"role": "user", "content": "Hi"}, + ], + ) + assert req.system == "Base.\n\nUse short answers." + assert "None" not in str(req.system) + assert [msg.role for msg in req.messages] == ["user"] + def test_tools_field_parses(self): req = AnthropicMessagesRequest( max_tokens = 100, @@ -157,6 +249,20 @@ class TestAnthropicMessagesToOpenAI: assert result[0] == {"role": "system", "content": "Be brief."} assert result[1] == {"role": "user", "content": "Hello"} + def test_top_level_system_request_translates_unchanged(self): + req = AnthropicMessagesRequest( + messages = [{"role": "user", "content": "Hello"}], + system = "Be brief.", + ) + result = anthropic_messages_to_openai( + [m.model_dump() for m in req.messages], + req.system, + ) + assert result == [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Hello"}, + ] + def test_system_as_block_list(self): system = [ {"type": "text", "text": "Be brief."}, @@ -312,10 +418,7 @@ class TestAnthropicMessagesToOpenAI: ] result = anthropic_messages_to_openai(msgs) parts = result[0]["content"] - assert parts[1] == { - "type": "image_url", - "image_url": {"url": "https://x/y.png"}, - } + assert parts[1] == {"type": "image_url", "image_url": {"url": "https://x/y.png"}} def test_image_only_user_message_emits_no_text_part(self): msgs = [ @@ -380,12 +483,7 @@ class TestAnthropicMessagesToOpenAI: ] result = anthropic_messages_to_openai(msgs) parts = result[0]["content"] - assert [p["type"] for p in parts] == [ - "text", - "image_url", - "text", - "image_url", - ] + assert [p["type"] for p in parts] == ["text", "image_url", "text", "image_url"] assert parts[0]["text"] == "before" assert parts[2]["text"] == "after" assert parts[1]["image_url"]["url"] == "data:image/png;base64,AA" @@ -463,15 +561,10 @@ class TestAnthropicToolsToOpenAI: enabled_tools = ["python"], ) - assert [tool["function"]["name"] for tool in result] == [ - "web_search", - "python", - ] + assert [tool["function"]["name"] for tool in result] == ["web_search", "python"] def test_pydantic_model_input(self): - tool = AnthropicTool( - name = "test", description = "desc", input_schema = {"type": "object"} - ) + tool = AnthropicTool(name = "test", description = "desc", input_schema = {"type": "object"}) result = anthropic_tools_to_openai([tool]) assert result[0]["function"]["name"] == "test" @@ -551,10 +644,52 @@ class TestAnthropicStreamEmitter: assert "tool_use" in events[1] assert "input_json_delta" in events[2] + def test_duplicate_tool_start_merges_into_open_tool_block(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + first_events = e.feed( + { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": "call_0", + "arguments": {}, + } + ) + second_events = e.feed( + { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": "call_0", + "arguments": {"code": ""}, + } + ) + + first_payloads = [json.loads(event.split("data: ")[1]) for event in first_events] + second_payloads = [json.loads(event.split("data: ")[1]) for event in second_events] + + tool_starts = [ + payload + for payload in first_payloads + second_payloads + if payload["type"] == "content_block_start" + and payload["content_block"]["type"] == "tool_use" + ] + assert len(tool_starts) == 1 + assert tool_starts[0]["content_block"]["id"].startswith("toolu_") + assert second_payloads == [ + { + "type": "content_block_delta", + "index": tool_starts[0]["index"], + "delta": { + "type": "input_json_delta", + "partial_json": json.dumps({"code": ""}), + }, + } + ] + def test_tool_end_closes_tool_opens_new_text_block(self): e = AnthropicStreamEmitter() e.start("msg_1", "m") - e.feed( + start_events = e.feed( { "type": "tool_start", "tool_name": "t", @@ -562,6 +697,13 @@ class TestAnthropicStreamEmitter: "arguments": {}, } ) + start_payload = next( + json.loads(event.split("data: ")[1]) + for event in start_events + if "content_block_start" in event + ) + tool_use_id = start_payload["content_block"]["id"] + assert tool_use_id.startswith("toolu_") events = e.feed( { "type": "tool_end", @@ -576,7 +718,7 @@ class TestAnthropicStreamEmitter: assert "tool_result" in events[1] parsed = json.loads(events[1].split("data: ")[1]) assert parsed["content"] == "done" - assert parsed["tool_use_id"] == "tc_1" + assert parsed["tool_use_id"] == tool_use_id assert "content_block_start" in events[2] assert '"type": "text"' in events[2] @@ -674,6 +816,44 @@ class TestAnthropicStreamEmitter: assert parsed["delta"]["text"] == "After tool" +# ===================================================================== +# Non-streaming tool response tests +# ===================================================================== + + +class TestAnthropicToolNonStreaming: + def test_duplicate_tool_start_replaces_provisional_tool_block(self): + def _run_gen(): + yield { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": "call_0", + "arguments": {}, + } + yield { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": "call_0", + "arguments": {"code": ""}, + } + yield { + "type": "tool_end", + "tool_name": "render_html", + "tool_call_id": "call_0", + "result": "Rendered HTML artifact.", + } + + response = asyncio.run(_anthropic_tool_non_streaming(_run_gen, "msg_1", "m")) + body = json.loads(response.body) + tool_blocks = [block for block in body["content"] if block["type"] == "tool_use"] + + assert len(tool_blocks) == 1 + assert tool_blocks[0]["type"] == "tool_use" + assert tool_blocks[0]["id"].startswith("toolu_") + assert tool_blocks[0]["name"] == "render_html" + assert tool_blocks[0]["input"] == {"code": ""} + + # ===================================================================== # Pass-through emitter tests (client-side tool execution path) # ===================================================================== @@ -739,7 +919,7 @@ class TestAnthropicPassthroughEmitter: parsed = self._parse(events[0]) assert parsed["type"] == "content_block_start" assert parsed["content_block"]["type"] == "tool_use" - assert parsed["content_block"]["id"] == "call_1" + assert parsed["content_block"]["id"].startswith("toolu_") assert parsed["content_block"]["name"] == "Bash" def test_tool_call_arguments_streamed_as_input_json_delta(self): @@ -768,26 +948,14 @@ class TestAnthropicPassthroughEmitter: events1 = e.feed_chunk( { "choices": [ - { - "delta": { - "tool_calls": [ - {"index": 0, "function": {"arguments": '{"cmd'}} - ] - } - } + {"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '{"cmd'}}]}} ] } ) events2 = e.feed_chunk( { "choices": [ - { - "delta": { - "tool_calls": [ - {"index": 0, "function": {"arguments": '": "ls"}'}} - ] - } - } + {"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '": "ls"}'}}]}} ] } ) @@ -956,7 +1124,102 @@ class TestAnthropicPassthroughEmitter: assert "content_block_start" in events[1] parsed = self._parse(events[1]) assert parsed["content_block"]["name"] == "Read" - assert parsed["content_block"]["id"] == "c2" + assert parsed["content_block"]["id"].startswith("toolu_") + + +class TestAnthropicPassthroughStreamAdapter: + class _Request: + async def is_disconnected(self): + return False + + @staticmethod + async def _collect(response): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk) + return chunks + + @staticmethod + def _payloads(lines, event_name): + prefix = f"event: {event_name}\n" + return [ + json.loads(line.split("data: ", 1)[1].strip()) + for line in lines + if line.startswith(prefix) + ] + + def test_stream_requests_usage_for_final_message_delta(self, monkeypatch): + import routes.inference as inf_mod + + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode()) + chunks = [ + {"choices": [{"delta": {"content": "hi"}}]}, + { + "choices": [], + "usage": { + "prompt_tokens": 2, + "completion_tokens": 4, + "total_tokens": 6, + }, + }, + ] + content = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + content += "data: [DONE]\n\n" + return httpx.Response( + 200, + content = content.encode(), + headers = {"content-type": "text/event-stream"}, + ) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def _client(*args, **kwargs): + return real_async_client( + transport = transport, + timeout = kwargs.get("timeout", 600), + ) + + monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client) + backend = SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + count_chat_tokens = lambda *args, **kwargs: 2, + ) + + async def run(): + response = await _anthropic_passthrough_stream( + self._Request(), + threading.Event(), + backend, + [{"role": "user", "content": "hi"}], + [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + 0.7, + 0.95, + 20, + 16, + "msg_1", + "test-model", + ) + return await self._collect(response) + + lines = asyncio.run(run()) + + assert captured["body"]["stream_options"] == {"include_usage": True} + message_delta = self._payloads(lines, "message_delta")[0] + assert message_delta["usage"]["input_tokens"] == 2 + assert message_delta["usage"]["output_tokens"] == 4 # ===================================================================== @@ -1069,14 +1332,14 @@ class TestAnthropicRequestedStudioTools: def test_bare_name_without_type_is_not_treated_as_server_tool(self): # Anthropic dispatches server tools by `type`; bare-name matching - # would let a malformed client tool (e.g. user forgot input_schema) - # silently flip the request into server-execution mode. + # would let a malformed client tool (missing input_schema) silently + # flip the request into server-execution mode. tools = [{"name": "python"}] assert _anthropic_requested_studio_tools(tools) == set() def test_client_tool_named_python_is_not_misclassified(self): - # input_schema is the client-tool discriminator; presence of it - # must prevent the name from being treated as a Studio alias. + # input_schema is the client-tool discriminator; its presence must + # prevent the name from being treated as a Studio alias. tools = [ { "name": "python", @@ -1110,27 +1373,23 @@ class TestAnthropicRequestedStudioTools: # ===================================================================== -class _PlainPathCalled(Exception): - pass - - -class _ToolPathCalled(Exception): - pass - - def _mock_backend(monkeypatch, **overrides): """Install a minimal stub backend on routes.inference. - Generation methods raise sentinel exceptions so the caller can assert - which path the route entered. + Generation methods record which path the route entered, then yield one + content event so the route can complete normally. """ import routes.inference as inf_mod + calls = [] + def _gen_plain(**kwargs): - raise _PlainPathCalled() + calls.append(("plain", kwargs)) + yield {"type": "content", "text": "ok"} def _gen_tools(**kwargs): - raise _ToolPathCalled() + calls.append(("tools", kwargs)) + yield {"type": "content", "text": "ok"} backend = SimpleNamespace( is_loaded = True, @@ -1139,6 +1398,7 @@ def _mock_backend(monkeypatch, **overrides): model_identifier = "test-model", generate_chat_completion = _gen_plain, generate_chat_completion_with_tools = _gen_tools, + calls = calls, ) backend.__dict__.update(overrides) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) @@ -1180,13 +1440,11 @@ class TestAnthropicMessagesToolRouting: assert exc.value.status_code == 400 assert "Mixing Anthropic server tools" in exc.value.detail - def test_mixed_rejected_when_client_tool_name_collides_with_server_alias( - self, monkeypatch - ): - # Regression: a client tool sharing a name with a mapped server - # tool (e.g. user defines their own "web_search") must still - # trigger the mixed-mode 400 — the post-name filter would - # otherwise drop the client tool and silently route to server-only. + def test_mixed_rejected_when_client_tool_name_collides_with_server_alias(self, monkeypatch): + # Regression: a client tool sharing a name with a mapped server tool + # (e.g. a custom "web_search") must still trigger the mixed-mode 400; + # otherwise the post-name filter drops the client tool and silently + # routes to server-only. _mock_backend(monkeypatch) payload = _basic_payload( tools = [ @@ -1213,10 +1471,10 @@ class TestAnthropicMessagesToolRouting: def test_client_tool_missing_name_rejected_with_400(self, monkeypatch): # Regression: AnthropicTool.name was relaxed to Optional for server - # tools, so a client-tool payload that has input_schema but omits - # `name` (e.g. typo) now parses successfully but would be silently - # dropped by anthropic_tools_to_openai, leaving the request with - # tool calling disabled. Reject at the boundary instead. + # tools, so a client-tool payload with input_schema but no `name` + # (typo) now parses but would be silently dropped by + # anthropic_tools_to_openai, leaving tool calling disabled. Reject at + # the boundary instead. _mock_backend(monkeypatch) payload = _basic_payload( tools = [{"input_schema": {"type": "object"}}], @@ -1230,7 +1488,7 @@ class TestAnthropicMessagesToolRouting: def test_client_tool_empty_name_rejected_with_400(self, monkeypatch): # Same silent-disable class as missing-name: `name: ""` passes the # isinstance check but is dropped by anthropic_tools_to_openai's - # `if not name` guard. Reject at the boundary so the typo surfaces. + # `if not name` guard. Reject at the boundary so the typo shows. _mock_backend(monkeypatch) payload = _basic_payload( tools = [{"name": "", "input_schema": {"type": "object"}}], @@ -1241,13 +1499,11 @@ class TestAnthropicMessagesToolRouting: assert exc.value.status_code == 400 assert "name" in exc.value.detail - def test_alias_named_client_tool_without_schema_rejected_with_400( - self, monkeypatch - ): - # Regression: a typo'd client tool whose name happens to collide - # with a Studio alias (e.g. user meant a custom "python" tool but - # forgot input_schema) must surface a 400, not silently switch - # the request into Studio's built-in python execution. + def test_alias_named_client_tool_without_schema_rejected_with_400(self, monkeypatch): + # Regression: a typo'd client tool whose name collides with a Studio + # alias (e.g. a custom "python" tool missing input_schema) must + # surface a 400, not silently switch into Studio's built-in python + # execution. _mock_backend(monkeypatch) payload = _basic_payload(tools = [{"name": "python"}]) @@ -1257,43 +1513,42 @@ class TestAnthropicMessagesToolRouting: assert "input_schema" in exc.value.detail def test_unrecognized_server_tool_accepted_as_noop(self, monkeypatch): - _mock_backend(monkeypatch) + backend = _mock_backend(monkeypatch) payload = _basic_payload( tools = [{"type": "code_execution_20250825", "name": "code_execution"}], ) - with pytest.raises(_PlainPathCalled): - _drive(anthropic_messages(payload, request = None, current_subject = "t")) + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert backend.calls[0][0] == "plain" def test_disable_tools_policy_overrides_server_tool_alias(self, monkeypatch): - # CLI `unsloth run --disable-tools` sets policy=False. A request - # carrying a Studio server-tool alias must NOT enter the agentic - # loop in that configuration. - _mock_backend(monkeypatch) + # CLI `unsloth run --disable-tools` sets policy=False. A request with + # a Studio server-tool alias must NOT enter the agentic loop then. + backend = _mock_backend(monkeypatch) set_tool_policy(False) payload = _basic_payload( tools = [{"type": "web_search_20250305", "name": "web_search"}], ) - with pytest.raises(_PlainPathCalled): - _drive(anthropic_messages(payload, request = None, current_subject = "t")) + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert backend.calls[0][0] == "plain" def test_server_tool_alias_enters_tool_path_when_policy_unset(self, monkeypatch): # Mirror of the previous test for the default (None) policy. - _mock_backend(monkeypatch) + backend = _mock_backend(monkeypatch) payload = _basic_payload( tools = [{"type": "web_search_20250305", "name": "web_search"}], ) - with pytest.raises(_ToolPathCalled): - _drive(anthropic_messages(payload, request = None, current_subject = "t")) + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert backend.calls[0][0] == "tools" def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch): - _mock_backend(monkeypatch) + backend = _mock_backend(monkeypatch) payload = _basic_payload( enable_tools = False, tools = [{"type": "web_search_20250305", "name": "web_search"}], ) - with pytest.raises(_PlainPathCalled): - _drive(anthropic_messages(payload, request = None, current_subject = "t")) + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert backend.calls[0][0] == "plain" diff --git a/studio/backend/tests/test_anthropic_thinking_translation.py b/studio/backend/tests/test_anthropic_thinking_translation.py index 14f261ae6b..02c4893735 100644 --- a/studio/backend/tests/test_anthropic_thinking_translation.py +++ b/studio/backend/tests/test_anthropic_thinking_translation.py @@ -1,24 +1,16 @@ # 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 Anthropic extended-thinking translation in -external_provider. +"""Unit tests for Anthropic extended-thinking translation in external_provider. Covers: -- Adaptive-mode request body nests effort under - ``output_config: {effort: ""}`` per the Messages API - reference (a top-level ``effort`` field 400s with - "effort: Extra inputs are not permitted"). -- Streaming SSE: ``content_block_delta`` with - ``delta.type == "thinking_delta"`` is translated into inline - ``...`` chat-completion chunks so the frontend's - reasoning-panel pipeline lifts it correctly. -- The ```` tag closes when the first ``text_delta`` arrives, - on ``content_block_stop``, on ``message_delta``, or on - ``message_stop``. -- Thinking is paired with ``temperature=1`` and no ``top_p`` / - ``top_k`` on the wire (Anthropic extended-thinking contract). +- Adaptive-mode body nests effort under ``output_config: {effort}`` (a + top-level ``effort`` field 400s). +- Streaming ``thinking_delta`` is translated into inline ``...`` + chunks for the frontend reasoning panel. +- The ```` tag closes on the first ``text_delta``, ``content_block_stop``, + ``message_delta``, or ``message_stop``. +- Thinking forces ``temperature=1`` with no ``top_p`` / ``top_k`` (contract). """ import asyncio @@ -113,9 +105,8 @@ def test_adaptive_thinking_body_uses_output_config_effort_shape(monkeypatch): # display=summarized is set explicitly so Opus 4.7 (which defaults to # "omitted") still emits thinking_delta events for the reasoning panel. assert body["thinking"] == {"type": "adaptive", "display": "summarized"} - # Documented shape: effort is nested under output_config. - # A top-level `effort` field produces a 400: - # "effort: Extra inputs are not permitted". + # Documented shape: effort is nested under output_config. A top-level + # `effort` field produces a 400: "effort: Extra inputs are not permitted". assert body["output_config"] == {"effort": "medium"} assert "effort" not in body # Extended-thinking contract: temperature=1, no top_p / top_k. @@ -258,10 +249,10 @@ def test_manual_thinking_body_uses_budget_tokens_on_4_5(monkeypatch): body = captured["body"] assert body["thinking"] == {"type": "enabled", "budget_tokens": 4096} # max_tokens must be strictly greater than budget_tokens; we shipped 1024 - # and budget is 4096, so the wrapper should bump max_tokens. + # and budget is 4096, so the wrapper must bump max_tokens. assert body["max_tokens"] > body["thinking"]["budget_tokens"] - # Manual-thinking path does not use output_config / effort — those are - # the adaptive-mode controls (Claude 4.6 / 4.7). + # Manual-thinking path does not use output_config / effort — those are the + # adaptive-mode controls (Claude 4.6 / 4.7). assert "effort" not in body assert "output_config" not in body @@ -338,8 +329,8 @@ def test_thinking_delta_wrapped_in_think_tags(monkeypatch): if isinstance(p, dict) and p["choices"][0]["delta"] ) - # Reasoning text should be wrapped in ..., followed by the - # answer text, and the stream should terminate with [DONE]. + # Reasoning text is wrapped in ..., then the answer text, and + # the stream terminates with [DONE]. assert "First I plan." in combined assert combined.endswith("Answer.") # signature_delta is intentionally dropped — no leaked signature text. @@ -350,9 +341,9 @@ def test_thinking_delta_wrapped_in_think_tags(monkeypatch): def test_thinking_only_turn_closes_tag_without_text_delta(monkeypatch): """display=omitted on Claude 4.7 emits a signature_delta and no text. - The open is still triggered by the (synthetic) thinking_delta; - we want content_block_stop to close it cleanly so the tag never leaks - into the next chunk.""" + The open is still triggered by the (synthetic) thinking_delta; we + want content_block_stop to close it cleanly so the tag never leaks into the + next chunk.""" def handler(request: httpx.Request) -> httpx.Response: events = [ diff --git a/studio/backend/tests/test_anthropic_tool_versions.py b/studio/backend/tests/test_anthropic_tool_versions.py index 608977ab07..fe2da26d73 100644 --- a/studio/backend/tests/test_anthropic_tool_versions.py +++ b/studio/backend/tests/test_anthropic_tool_versions.py @@ -1,35 +1,13 @@ # 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 per-model Anthropic server-side tool-version dispatch -helpers in ``core.inference.external_provider``. +"""Tests for the per-model Anthropic tool-version dispatch helpers in +``core.inference.external_provider``. -Anthropic ships date-pinned tool versions per model family. The newer -``_20260209`` web_search / web_fetch and ``_20260120`` code_execution -variants only run on a subset of models; sending them to an older -model returns a 400 from upstream, and sending the older -``_20250305`` / ``_20250910`` / ``_20250825`` variants to a newer -model misses dynamic filtering and the free-when-paired pricing. The -helpers below decide which version goes out per model; this test pins -the dispatch matrix so future model launches keep working without -silently regressing the newer-version path. - -Covers: -- ``_anthropic_web_search_version`` / ``_anthropic_web_fetch_version`` - pick ``_20260209`` for Opus 4.6+, Opus 4.7, Sonnet 4.6 and fall back - to ``_20250305`` / ``_20250910`` for everything else (4.5 family, - Haiku 4.5, 4.1, 4.0). -- ``_anthropic_code_execution_version`` picks ``_20260120`` for the - Opus 4.5+ / Sonnet 4.5+ / Opus 4.7 / Sonnet 4.6 family and falls back - to ``_20250825`` everywhere else (Haiku 4.5, 4.1, 4.0). -- ``_stream_anthropic`` body integration: when ``enabled_tools= - ["web_search", "code_execution"]`` is set on Opus 4.7, the outbound - body carries the newer pinned versions; the same payload on Haiku - 4.5 falls back to the legacy versions. -- The ``anthropic-beta: code-execution-2025-08-25`` header is sent - unchanged for both code-execution variants (no header rev needed). -""" +Anthropic ships date-pinned tool versions per model family; sending the +wrong-dated variant to a model 400s upstream. Pins the dispatch matrix for +web_search/web_fetch/code_execution helpers, the ``_stream_anthropic`` body +integration, and the unchanged code-execution beta header.""" import asyncio import json @@ -165,13 +143,8 @@ def test_outbound_body_uses_new_versions_on_opus_4_7(monkeypatch): assert "code_execution_20260120" in tool_types assert "web_search_20250305" not in tool_types assert "code_execution_20250825" not in tool_types - # Beta header for code execution stays on the existing flag for - # both _20250825 and _20260120; the API uses one header to gate - # the feature, not the date. - assert "code-execution-2025-08-25" in captured["headers"].get( - "anthropic-beta", - "", - ) + # One beta header gates both _20250825 and _20260120. + assert "code-execution-2025-08-25" in captured["headers"].get("anthropic-beta", "") def test_outbound_body_falls_back_on_haiku_4_5(monkeypatch): diff --git a/studio/backend/tests/test_anthropic_web_fetch.py b/studio/backend/tests/test_anthropic_web_fetch.py index da10d679eb..bf24175256 100644 --- a/studio/backend/tests/test_anthropic_web_fetch.py +++ b/studio/backend/tests/test_anthropic_web_fetch.py @@ -1,12 +1,11 @@ # 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 `web_fetch_20250910` / `web_fetch_20260209` +"""Unit tests for Anthropic's `web_fetch_20250910` / `web_fetch_20260209` translation in ``_stream_anthropic``. Covers request body emission -(version picked by ``_anthropic_web_fetch_version``: ``_20260209`` for -Opus 4.6/4.7 + Sonnet 4.6, ``_20250910`` otherwise), combined tool -requests, off-by-default behavior, and SSE translation of success and +(version from ``_anthropic_web_fetch_version``: ``_20260209`` for Opus +4.6/4.7 + Sonnet 4.6, else ``_20250910``), combined tool requests, +off-by-default behavior, and SSE translation of success and ``url_not_accessible`` error paths into ``tool_start`` / ``tool_end``. """ @@ -103,12 +102,8 @@ def test_web_fetch_tool_appended_to_request_body(monkeypatch): body = captured["body"] tools = body.get("tools") or [] - # claude-opus-4-7 routes web_fetch to _20260209 (dynamic filtering). - assert { - "type": "web_fetch_20260209", - "name": "web_fetch", - "max_uses": 5, - } in tools + # claude-opus-4-7 routes web_fetch to _20260209. + assert {"type": "web_fetch_20260209", "name": "web_fetch", "max_uses": 5} in tools # web_fetch is GA; no beta header is required. assert "web-fetch" not in captured["headers"].get("anthropic-beta", "") @@ -144,13 +139,13 @@ def test_web_fetch_combined_with_web_search_and_code_execution(monkeypatch): tools = captured["body"].get("tools") or [] tool_types = [t.get("type") for t in tools] - # claude-opus-4-7 routes web_search and web_fetch to _20260209 - # and code_execution to _20260120 (per PR 5679 dispatch). + # claude-opus-4-7 routes web_search/web_fetch to _20260209 and + # code_execution to _20260120 (per PR 5679 dispatch). assert "web_search_20260209" in tool_types, tool_types assert "web_fetch_20260209" in tool_types, tool_types assert "code_execution_20260120" in tool_types, tool_types - # Code-execution still adds its beta flag; web_fetch must not - # have accidentally stripped it. + # Code-execution still adds its beta flag; web_fetch must not have + # stripped it. assert "code-execution-2025-08-25" in captured["headers"].get("anthropic-beta", "") @@ -183,9 +178,7 @@ def test_no_web_fetch_tool_when_pill_off(monkeypatch): _drive(run()) tools = captured["body"].get("tools") or [] - assert all( - t.get("type") not in ("web_fetch_20250910", "web_fetch_20260209") for t in tools - ) + assert all(t.get("type") not in ("web_fetch_20250910", "web_fetch_20260209") for t in tools) # ── SSE translation ───────────────────────────────────────────────── @@ -253,9 +246,7 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch): client = _make_client() return await _collect( client._stream_anthropic( - messages = [ - {"role": "user", "content": "Fetch https://example.com/article"} - ], + messages = [{"role": "user", "content": "Fetch https://example.com/article"}], model = "claude-opus-4-7", temperature = 0.7, top_p = 0.95, @@ -271,7 +262,9 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch): assert start["type"] == "tool_start" assert start["tool_name"] == "web_fetch" assert start["tool_call_id"] == "srvtoolu_wf1" - assert start["arguments"] == {"url": "https://example.com/article"} + # `_server_tool: True` marks this a provider-side synthetic tool card + # for the frontend's history serializer. + assert start["arguments"] == {"url": "https://example.com/article", "_server_tool": True} assert end["type"] == "tool_end" assert end["tool_call_id"] == "srvtoolu_wf1" # The source pill uses Title / URL / snippet as parseSourcesFromResult expects. @@ -351,9 +344,10 @@ def test_web_fetch_error_renders_error_code(monkeypatch): def _finish_reasons(lines: list[str]) -> list: - """Return non-null finish_reason fields from each chat.completion.chunk. - Mid-stream content deltas carry ``finish_reason: None`` and are skipped - (the refusal path emits a notice delta before the content_filter chunk).""" + """Non-null finish_reason fields from each chat.completion.chunk. + Mid-stream content deltas carry ``finish_reason: None`` and are + skipped (refusal emits a notice delta before the content_filter + chunk).""" out: list = [] for line in lines: if not line.startswith("data:"): @@ -375,12 +369,10 @@ def _finish_reasons(lines: list[str]) -> list: def test_pause_turn_does_not_emit_finish_reason_chunk(monkeypatch): - # `pause_turn` is what Anthropic emits when a long server-tool turn - # (typically web_search / web_fetch) pauses and will resume on the - # next request. Treating it as finish_reason="stop" makes the - # OpenAI-formatted client truncate the rendered assistant message. - # The adapter must skip the chunk so the stream ends cleanly with - # [DONE] and no terminal finish_reason. + # Anthropic emits `pause_turn` when a long server-tool turn pauses and + # resumes next request. Mapping it to finish_reason="stop" truncates the + # OpenAI client's message, so the adapter must skip the chunk and end + # cleanly with [DONE] and no terminal finish_reason. sse_events = [ {"type": "message_start", "message": {"usage": {}}}, { @@ -414,15 +406,14 @@ def test_pause_turn_does_not_emit_finish_reason_chunk(monkeypatch): ) lines = _drive(run()) - # No finish_reason chunk for pause_turn -- the only completion - # signal is the [DONE] line. + # No finish_reason chunk for pause_turn -- only [DONE] signals + # completion. assert _finish_reasons(lines) == [], lines assert any(line.strip() == "data: [DONE]" for line in lines), lines def test_end_turn_still_emits_stop_finish_reason(monkeypatch): - # Sanity: the pause_turn -> None mapping must not regress normal - # end_turn handling. + # Sanity: pause_turn -> None mapping must not regress end_turn. sse_events = [ {"type": "message_start", "message": {"usage": {}}}, { @@ -495,12 +486,9 @@ def test_refusal_maps_to_content_filter(monkeypatch): def test_web_fetch_titleless_document_falls_back_to_url(monkeypatch): - # Anthropic may omit `document.title` on pages where the HTML - # provides nothing usable. Without a fallback the formatter would - # emit `URL: ...\nSnippet: ...` only, and the frontend's - # parseSourcesFromResult skips entries that lack a `Title:` line, - # so the source pill silently disappears. Verify the formatter - # mirrors the web_search behaviour and falls back to the URL. + # Anthropic may omit `document.title`. Without a fallback the formatter + # emits no `Title:` line, and the frontend's parseSourcesFromResult drops + # those entries (source pill disappears). Verify it falls back to the URL. sse_events = [ {"type": "message_start", "message": {"usage": {}}}, { diff --git a/studio/backend/tests/test_audio_token_detection.py b/studio/backend/tests/test_audio_token_detection.py new file mode 100644 index 0000000000..02d4790519 --- /dev/null +++ b/studio/backend/tests/test_audio_token_detection.py @@ -0,0 +1,43 @@ +# 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 tokenizer-based audio_type detection, covering Gemma 3n +() and Gemma 4 (<|audio|>) audio-input tokens.""" + +from __future__ import annotations + +from utils.models.model_config import _AUDIO_TOKEN_PATTERNS, is_audio_input_type + + +def _classify(tokens: list[str]) -> str | None: + """Mirror _check_token_patterns: first match in dict order wins.""" + for audio_type, check in _AUDIO_TOKEN_PATTERNS.items(): + if check(tokens): + return audio_type + return None + + +def test_gemma3n_audio_soft_token_is_audio_vlm(): + assert _classify(["", "", ""]) == "audio_vlm" + + +def test_gemma4_pipe_audio_token_is_audio_vlm(): + # Gemma 4 uses <|audio|> (and <|image|>) instead of *_soft_token. + assert _classify(["", "<|image|>", "<|audio|>"]) == "audio_vlm" + + +def test_csm_uppercase_audio_not_classified_as_audio_vlm(): + # csm uses uppercase <|AUDIO|> + <|audio_eos|>; must stay csm, not audio_vlm. + tokens = ["<|AUDIO|>", "<|audio_eos|>"] + assert _classify(tokens) == "csm" + + +def test_audio_vlm_and_whisper_accept_audio_input(): + assert is_audio_input_type("audio_vlm") is True + assert is_audio_input_type("whisper") is True + assert is_audio_input_type("snac") is False + assert is_audio_input_type(None) is False + + +def test_non_audio_tokens_classify_none(): + assert _classify(["", "", ""]) is None diff --git a/studio/backend/tests/test_browse_folders_route.py b/studio/backend/tests/test_browse_folders_route.py index 19a83987d3..3a607e6b10 100644 --- a/studio/backend/tests/test_browse_folders_route.py +++ b/studio/backend/tests/test_browse_folders_route.py @@ -9,8 +9,7 @@ from pathlib import Path import pytest from fastapi import HTTPException -# Keep this test runnable in lightweight environments where optional logging -# deps are not installed. +# Keep runnable in lightweight environments lacking optional logging deps. if "structlog" not in sys.modules: class _DummyLogger: diff --git a/studio/backend/tests/test_cache_case_resolution.py b/studio/backend/tests/test_cache_case_resolution.py index 60963b4f7d..22e961891e 100644 --- a/studio/backend/tests/test_cache_case_resolution.py +++ b/studio/backend/tests/test_cache_case_resolution.py @@ -5,8 +5,7 @@ from pathlib import Path import sys import types -# Keep this test runnable in lightweight environments where optional logging -# deps are not installed. +# Stub structlog so this test runs where the optional dep is absent. if "structlog" not in sys.modules: class _DummyLogger: @@ -111,10 +110,10 @@ def test_resolve_cached_repo_id_case_late_cache_population(tmp_path, monkeypatch first = resolve_cached_repo_id_case("org/model") assert first == "org/model" - # Simulate cache being populated after first miss (e.g. another code path/download). + # Cache populated after first miss (e.g. another code path/download). _mk_cache_repo(tmp_path, "Org/Model") second = resolve_cached_repo_id_case("org/model") - # Desired behavior: second lookup should pick up the now-existing variant. + # Second lookup should pick up the now-existing variant. assert second == "Org/Model" diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 05aae8fb75..20aa37a2b4 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -7,8 +7,7 @@ import types from pathlib import Path from types import SimpleNamespace -# Keep this test runnable in lightweight environments where optional logging -# deps are not installed. +# Keep this test runnable without optional logging deps. if "structlog" not in sys.modules: class _DummyLogger: @@ -66,9 +65,7 @@ def test_iter_gguf_paths_matches_extension_case_insensitively(tmp_path): assert result == ["Q4_K_M.gguf", "Q8_0.GGUF"] -def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf( - monkeypatch, tmp_path -): +def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monkeypatch, tmp_path): repo = _repo( "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive", [_file("Q4_K_M.gguf", 5_000), _file("README.md", 10)], @@ -130,9 +127,7 @@ def test_list_cached_gguf_skips_repos_without_positive_gguf_size(monkeypatch, tm assert result["cached"] == [] -def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans( - monkeypatch, tmp_path -): +def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans(monkeypatch, tmp_path): smaller = _repo( "Org/Dupe", [_file("Q4_K_M.gguf", 2_000)], @@ -193,9 +188,7 @@ def test_list_cached_gguf_dedupes_shared_blobs_across_revisions(monkeypatch, tmp ] -def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist( - monkeypatch, tmp_path -): +def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(monkeypatch, tmp_path): mixed = _repo( "Org/MixedRepo", [ @@ -216,11 +209,8 @@ def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist( assert result["cached"] == [] -def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors( - monkeypatch, tmp_path -): - """Mirror of the _skips_ test: the mixed repo should still surface in - cached-gguf so the picker can show it as a GGUF download.""" +def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypatch, tmp_path): + """Mixed repo still surfaces in cached-gguf as a GGUF download.""" mixed = _repo( "Org/MixedRepo", [ @@ -248,9 +238,8 @@ def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors( def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path): - """A partial/interrupted GGUF download has ``size_on_disk = None``. The - route must treat the unknown bytes as zero instead of raising TypeError - out of ``sum()`` and wiping the entire response.""" + """``size_on_disk = None`` (partial download) is treated as zero, not a + TypeError from ``sum()`` that wipes the response.""" partial = _repo( "Org/PartialDownload", [_file("Q4_K_M.gguf", None), _file("Q6_K.gguf", 5_000)], @@ -274,11 +263,8 @@ def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path): ] -def test_list_cached_gguf_skips_malformed_repo_without_wiping_response( - monkeypatch, tmp_path -): - """One repo raising during classification must not poison the response - for every other repo in the scan.""" +def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(monkeypatch, tmp_path): + """One repo raising during classification must not poison the response.""" class _ExplodingRepo: repo_id = "Org/Broken" @@ -313,9 +299,8 @@ def test_list_cached_gguf_skips_malformed_repo_without_wiping_response( def test_list_cached_gguf_skips_repo_with_only_mmproj_gguf(monkeypatch, tmp_path): - """A repo whose only ``.gguf`` artifact is an mmproj vision adapter - must not be classified as a GGUF repo: the variant selector filters - mmproj out and the picker would otherwise show zero variants.""" + """A repo whose only ``.gguf`` is an mmproj vision adapter is not a GGUF + repo: mmproj is filtered out, leaving zero variants.""" mmproj_only = _repo( "Org/MmprojOnly", [ @@ -337,9 +322,8 @@ def test_list_cached_gguf_skips_repo_with_only_mmproj_gguf(monkeypatch, tmp_path def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp_path): - """Mirror of the cached-gguf skip: a safetensors repo with an - auxiliary mmproj vision adapter must still surface in cached-models - so the user can load it as a normal model.""" + """A safetensors repo with an auxiliary mmproj adapter still surfaces in + cached-models as a normal model.""" mmproj_aux = _repo( "Org/MmprojAux", [ @@ -357,21 +341,12 @@ def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) - assert result["cached"] == [ - { - "repo_id": "Org/MmprojAux", - "size_bytes": 15_000, - } - ] + assert result["cached"] == [{"repo_id": "Org/MmprojAux", "size_bytes": 15_000}] -def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj( - monkeypatch, tmp_path -): - """A vision-capable GGUF repo (main weight + mmproj adapter) is still - a GGUF repo. The reported size is the main weight size; mmproj is - excluded from the GGUF-size accounting because it is filtered out at - classification time.""" +def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeypatch, tmp_path): + """A vision GGUF repo (main weight + mmproj) is a GGUF repo; reported size + is the main weight only, since mmproj is filtered at classification.""" vision_repo = _repo( "Org/VisionGguf", [ diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index 9337544638..e3aeb4eb22 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -3,6 +3,7 @@ import asyncio import os +import re import sys import pytest @@ -56,6 +57,68 @@ def test_replace_thread_messages_rejects_body_thread_mismatch(monkeypatch): assert called is False +# --------------------------------------------------------------------------- +# /api/chat/settings +# --------------------------------------------------------------------------- + + +def test_chat_settings_payload_accepts_fast_mode_presets(): + payload = chat_history.ChatSettingsPayload.model_validate( + { + "inferenceParams": {"fastMode": False}, + "customPresets": [ + { + "name": "Fast Opus", + "params": { + "temperature": 0.6, + "topP": 0.95, + "topK": 20, + "minP": 0.01, + "repetitionPenalty": 1.0, + "presencePenalty": 0.0, + "maxTokens": 8192, + "systemPrompt": "", + "trustRemoteCode": False, + "fastMode": True, + }, + }, + ], + } + ) + + dumped = payload.model_dump(exclude_unset = True) + assert dumped["inferenceParams"]["fastMode"] is False + assert dumped["customPresets"][0]["params"]["fastMode"] is True + + +def test_chat_inference_settings_covers_frontend_persisted_fields(): + # Drift guard: every InferenceParams field the UI persists (all but + # checkpoint) must exist on ChatInferenceSettings, else extra="forbid" + # 400s PUT /api/chat/settings on the next added field (issue #5862). + runtime_ts = os.path.join( + _backend, + "..", + "frontend", + "src", + "features", + "chat", + "types", + "runtime.ts", + ) + if not os.path.exists(runtime_ts): + pytest.skip("frontend runtime.ts not present") + + with open(runtime_ts, encoding = "utf-8") as fh: + block = re.search(r"interface InferenceParams \{(.*?)\n\}", fh.read(), re.DOTALL) + assert block, "InferenceParams interface not found in runtime.ts" + persisted = set(re.findall(r"^\s*(\w+)\??:", block.group(1), re.M)) - {"checkpoint"} + + backend = set(chat_history.ChatInferenceSettings.model_fields) + assert persisted == backend, ( + f"schema drift: frontend-only {persisted - backend}, " f"backend-only {backend - persisted}" + ) + + # --------------------------------------------------------------------------- # /api/chat/import-ledger # --------------------------------------------------------------------------- @@ -102,7 +165,6 @@ def test_record_import_ledger_returns_accepted_and_inserted(monkeypatch): def test_record_import_ledger_rejects_oversize_payload(): from pydantic import ValidationError - with pytest.raises(ValidationError): chat_history.ChatImportLedgerRecordRequest( threadIds = [f"id-{i}" for i in range(10_001)], diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index 123dbf1b96..bc74ec172a 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -1,18 +1,52 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import os +import platform +import shutil import threading +import uuid +from pathlib import Path import pytest from storage import studio_db -def _reset_studio_db(tmp_path, monkeypatch): +def _reset_studio_db( + tmp_path, + monkeypatch, + projects_home = None, +): monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setenv( + "UNSLOTH_STUDIO_PROJECTS_HOME", + str(projects_home if projects_home is not None else tmp_path / "Projects"), + ) monkeypatch.setattr(studio_db, "_schema_ready", False) +@pytest.fixture +def workspace_projects_home(tmp_path): + """Projects root outside the platform delete denylist. + + macOS tmp_path resolves under /private/tmp, which the delete guard refuses; + only the denied case falls back to a home subdir. + """ + candidate = tmp_path / "Projects" + resolved = str(candidate.resolve()) + check = os.path.normcase(resolved) if platform.system() == "Windows" else resolved + denied = studio_db._denied_path_prefixes() + if any(check == p or check.startswith(p + os.sep) for p in denied): + candidate = Path.home() / ".unsloth-studio-tests" / uuid.uuid4().hex + candidate.mkdir(parents = True, exist_ok = True) + try: + yield candidate + finally: + if ".unsloth-studio-tests" in candidate.parts: + shutil.rmtree(candidate, ignore_errors = True) + + def _thread(thread_id: str = "thread-1") -> dict: return { "id": thread_id, @@ -41,6 +75,17 @@ def _message( } +def _project(project_id: str = "project-1") -> dict: + return { + "id": project_id, + "name": "Research", + "instructions": "Use terse answers.", + "archived": False, + "createdAt": 1_700_000_000_000, + "updatedAt": 1_700_000_000_000, + } + + def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) studio_db.upsert_chat_thread(_thread()) @@ -63,6 +108,50 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch): assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}] +def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + project = studio_db.upsert_chat_project(_project()) + assert project["rootPath"].startswith(str(tmp_path / "Projects")) + assert (tmp_path / "Projects" / "Research-project").exists() + assert (tmp_path / "Projects" / "Research-project" / "sandbox").is_dir() + assert not (tmp_path / "Projects" / "Research-project" / "chats").exists() + assert not (tmp_path / "Projects" / "Research-project" / "files").exists() + assert not (tmp_path / "Projects" / "Research-project" / "exports").exists() + studio_db.upsert_chat_thread({**_thread(), "projectId": "project-1"}) + studio_db.upsert_chat_message(_message("msg-1", 1, "delete with project")) + + [thread] = studio_db.list_chat_threads(project_id = "project-1") + assert thread["projectId"] == "project-1" + + deleted = studio_db.delete_chat_project("project-1") + + assert deleted is not None + assert deleted["id"] == "project-1" + assert studio_db.get_chat_project("project-1") is None + assert studio_db.list_chat_threads(project_id = "project-1") == [] + assert studio_db.get_chat_thread("thread-1") is None + assert studio_db.list_chat_messages("thread-1") == [] + assert (tmp_path / "Projects" / "Research-project").exists() + + +def test_chat_project_delete_files_removes_workspace( + tmp_path, monkeypatch, workspace_projects_home +): + _reset_studio_db(tmp_path, monkeypatch, projects_home = workspace_projects_home) + project = studio_db.upsert_chat_project(_project()) + # Derive root from the created project so it tracks the projects home. + root = Path(project["rootPath"]) + marker = root / "sandbox" / "marker.txt" + marker.write_text("created by code execution", encoding = "utf-8") + + deleted = studio_db.delete_chat_project(project["id"], delete_files = True) + + assert deleted is not None + assert deleted["rootPath"] == project["rootPath"] + assert not root.exists() + assert studio_db.get_chat_project(project["id"]) is None + + def test_sync_chat_messages_prunes_when_requested(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) studio_db.upsert_chat_thread(_thread()) @@ -142,23 +231,16 @@ def test_settings_merge_atomic_under_concurrency(tmp_path, monkeypatch): def test_settings_merge_preserves_nested_keys(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) - studio_db.upsert_chat_settings_merge( - {"inferenceParams": {"temperature": 0.5, "topP": 0.8}} - ) + studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.5, "topP": 0.8}}) studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.9}}) params = studio_db.list_chat_settings()["inferenceParams"] assert params == {"temperature": 0.9, "topP": 0.8} -def test_settings_merge_quarantines_corrupt_json_and_rejects_partial_patch( - tmp_path, - monkeypatch, -): +def test_settings_merge_quarantines_corrupt_json_and_rejects_partial_patch(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) - studio_db.upsert_chat_settings_merge( - {"inferenceParams": {"temperature": 0.5, "topP": 0.8}} - ) + studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.5, "topP": 0.8}}) conn = studio_db.get_connection() try: conn.execute( @@ -206,9 +288,7 @@ def test_settings_merge_replaces_corrupt_scalar_after_quarantine(tmp_path, monke assert settings["autoTitle"] is True conn = studio_db.get_connection() try: - quarantined = conn.execute( - "SELECT key, reason FROM chat_settings_quarantine" - ).fetchall() + quarantined = conn.execute("SELECT key, reason FROM chat_settings_quarantine").fetchall() finally: conn.close() assert [(row["key"], row["reason"]) for row in quarantined] == [ @@ -264,11 +344,7 @@ def test_legacy_imports_records_and_lists(tmp_path, monkeypatch): ) assert accepted == 3 assert inserted == 3 - assert set(studio_db.list_chat_legacy_imports()) == { - "legacy-a", - "legacy-b", - "legacy-c", - } + assert set(studio_db.list_chat_legacy_imports()) == {"legacy-a", "legacy-b", "legacy-c"} def test_legacy_imports_is_idempotent(tmp_path, monkeypatch): @@ -282,11 +358,7 @@ def test_legacy_imports_is_idempotent(tmp_path, monkeypatch): assert (accepted1, inserted1) == (2, 2) # legacy-b is already in the ledger, only legacy-c is genuinely new. assert (accepted2, inserted2) == (2, 1) - assert set(studio_db.list_chat_legacy_imports()) == { - "legacy-a", - "legacy-b", - "legacy-c", - } + assert set(studio_db.list_chat_legacy_imports()) == {"legacy-a", "legacy-b", "legacy-c"} def test_legacy_imports_dedups_input(tmp_path, monkeypatch): @@ -294,8 +366,8 @@ def test_legacy_imports_dedups_input(tmp_path, monkeypatch): accepted, inserted = studio_db.upsert_chat_legacy_imports( ["x", "x", "y", "x"], ) - # accepted is the deduped non-empty input size; inserted is the rows - # actually new in the ledger after ON CONFLICT DO NOTHING. + # accepted is the deduped non-empty input size; inserted is the rows newly + # added to the ledger after ON CONFLICT DO NOTHING. assert accepted == 2 assert inserted == 2 assert set(studio_db.list_chat_legacy_imports()) == {"x", "y"} diff --git a/studio/backend/tests/test_cleanup_cancelled_checkpoints.py b/studio/backend/tests/test_cleanup_cancelled_checkpoints.py index 0d09f027cf..3799a4f91b 100644 --- a/studio/backend/tests/test_cleanup_cancelled_checkpoints.py +++ b/studio/backend/tests/test_cleanup_cancelled_checkpoints.py @@ -16,11 +16,10 @@ if str(_BACKEND_ROOT) not in sys.path: @pytest.fixture def outputs_setup(tmp_path, monkeypatch): - """Point outputs_root() at a temp dir so cleanup is allowed to run on it. + """Point outputs_root() at a temp dir so cleanup may run on it. - The training module binds ``outputs_root`` at import time - (``from utils.paths import outputs_root``), so we have to patch - the symbol on the importer module, not on storage_roots. + training binds ``outputs_root`` at import time, so patch the symbol on + the importer module, not on storage_roots. """ from core.training import training as training_mod @@ -36,8 +35,8 @@ def _mk_dir(parent: Path, name: str) -> Path: def test_completed_checkpoints_are_preserved(outputs_setup): - """The big regression: prior to this fix, every completed - checkpoint-N/ was rmtree'd on Cancel, destroying resume points.""" + """Regression: completed checkpoint-N/ used to be rmtree'd on Cancel, + destroying resume points.""" from core.training.training import _cleanup_cancelled_checkpoints out = outputs_setup / "run-1" @@ -130,7 +129,6 @@ def test_symlinked_output_dir_skipped(outputs_setup): def test_missing_output_dir_is_noop(outputs_setup): from core.training.training import _cleanup_cancelled_checkpoints - _cleanup_cancelled_checkpoints(outputs_setup / "does-not-exist") # Should not raise; nothing to assert beyond non-failure. diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py new file mode 100644 index 0000000000..873547631d --- /dev/null +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -0,0 +1,479 @@ +# 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 the Cloudflare quick-tunnel helper and run.py wiring. + +cloudflare_tunnel.py is stdlib-only (storage_roots is imported lazily), so it +loads via spec_from_file_location without the studio venv. run.py defaults are +checked by AST so we never import its heavy deps (uvicorn/structlog). +""" + +import ast +import importlib.util +import io +import sys +import tarfile +import types +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +_CT_PY = _BACKEND / "cloudflare_tunnel.py" +_RUN_PY = _BACKEND / "run.py" + + +def _load_ct(): + spec = importlib.util.spec_from_file_location("cloudflare_tunnel", _CT_PY) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +ct = _load_ct() + + +# ── URL parsing ────────────────────────────────────────────────────── + + +def test_url_regex_extracts_and_ignores_noise(): + blob = ( + "2026-06-11T10:00:00Z INF Thank you for trying Cloudflare Tunnel.\n" + "2026-06-11T10:00:01Z INF Requesting new quick Tunnel on trycloudflare.com...\n" + "2026-06-11T10:00:01Z INF | https://setting-democracy-gathering.trycloudflare.com |\n" + "2026-06-11T10:00:02Z INF Registered tunnel connection https://not-the-url.example.com\n" + ) + m = ct._URL_RE.search(blob) + assert m is not None + assert m.group(0) == "https://setting-democracy-gathering.trycloudflare.com" + + +def test_url_regex_no_match_on_unrelated(): + assert ct._URL_RE.search("INF connecting to https://api.cloudflare.com/v4") is None + + +# ── asset mapping ──────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "system,machine,expected", + [ + ("Linux", "x86_64", ("cloudflared-linux-amd64", False)), + ("Linux", "aarch64", ("cloudflared-linux-arm64", False)), + ("Darwin", "arm64", ("cloudflared-darwin-arm64.tgz", True)), + ("Darwin", "x86_64", ("cloudflared-darwin-amd64.tgz", True)), + ("Windows", "AMD64", ("cloudflared-windows-amd64.exe", False)), + ("Windows", "x86", ("cloudflared-windows-386.exe", False)), + ("Linux", "mips", None), + ("Plan9", "x86_64", None), + ], +) +def test_asset_name(monkeypatch, system, machine, expected): + monkeypatch.setattr(ct.platform, "system", lambda: system) + monkeypatch.setattr(ct.platform, "machine", lambda: machine) + assert ct._asset_name() == expected + + +# ── binary discovery ───────────────────────────────────────────────── + + +def test_find_cloudflared_prefers_path(monkeypatch): + monkeypatch.setattr(ct.shutil, "which", lambda name: "/usr/local/bin/cloudflared") + assert ct.find_cloudflared() == "/usr/local/bin/cloudflared" + + +def test_find_cloudflared_falls_back_to_cache(monkeypatch, tmp_path): + cached = tmp_path / "cloudflared" + cached.write_text("#!/bin/sh\n") + cached.chmod(0o755) + monkeypatch.setattr(ct.shutil, "which", lambda name: None) + monkeypatch.setattr(ct, "_cache_path", lambda: cached) + assert ct.find_cloudflared() == str(cached) + + +def test_find_cloudflared_none_when_missing(monkeypatch, tmp_path): + monkeypatch.setattr(ct.shutil, "which", lambda name: None) + monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "absent") + assert ct.find_cloudflared() is None + + +# ── ensure / download ──────────────────────────────────────────────── + + +def test_ensure_downloads_and_chmods_when_missing(monkeypatch, tmp_path): + cached = tmp_path / "cloudflared" + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-linux-amd64", False)) + monkeypatch.setattr(ct, "_cache_path", lambda: cached) + + def fake_download(url, dest): + assert url.endswith("/cloudflared-linux-amd64") + dest.write_bytes(b"ELF-ish") + return True + + monkeypatch.setattr(ct, "_download", fake_download) + monkeypatch.setattr(ct.sys, "platform", "linux") + path = ct.ensure_cloudflared() + assert path == str(cached) + assert cached.exists() + assert cached.stat().st_mode & 0o111 # executable bit set + + +def test_ensure_returns_none_on_download_failure(monkeypatch, tmp_path): + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-linux-amd64", False)) + monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "cloudflared") + monkeypatch.setattr(ct, "_download", lambda url, dest: False) + assert ct.ensure_cloudflared() is None + + +def test_ensure_returns_none_for_unsupported_arch(monkeypatch, tmp_path): + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: None) + monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "cloudflared") + assert ct.ensure_cloudflared() is None + + +def test_download_sets_user_agent(monkeypatch, tmp_path): + import urllib.request + + captured = {} + + class _Resp: + _sent = False + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, n = -1): + if self._sent: + return b"" + self._sent = True + return b"payload" + + def fake_urlopen(req, timeout = None): + captured["ua"] = req.get_header("User-agent") + return _Resp() + + monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen) + dest = tmp_path / "cloudflared" + assert ct._download("https://github.com/cloudflare/cloudflared/x", dest) is True + assert captured["ua"] == "unsloth-studio" # GitHub CDN 403s the default UA + assert dest.read_bytes() == b"payload" + + +# ── cross-platform: Windows (.exe), macOS (.tgz) ───────────────────── + + +def test_cache_path_uses_exe_on_windows(monkeypatch, tmp_path): + import types + + fake_sr = types.ModuleType("utils.paths.storage_roots") + fake_sr.studio_bin_root = lambda: tmp_path + monkeypatch.setitem(sys.modules, "utils.paths.storage_roots", fake_sr) + monkeypatch.setattr(ct.sys, "platform", "win32") + assert ct._cache_path() == tmp_path / "cloudflared.exe" + + +def test_ensure_windows_downloads_exe(monkeypatch, tmp_path): + cached = tmp_path / "cloudflared.exe" + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-windows-amd64.exe", False)) + monkeypatch.setattr(ct, "_cache_path", lambda: cached) + monkeypatch.setattr(ct.sys, "platform", "win32") + + def fake_download(url, dest): + assert url.endswith("/cloudflared-windows-amd64.exe") + dest.write_bytes(b"MZ") # PE header magic + return True + + monkeypatch.setattr(ct, "_download", fake_download) + # chmod is skipped on Windows; would raise on a path that does not exist yet. + monkeypatch.setattr(ct.os, "chmod", lambda *a, **k: pytest.fail("chmod called on win32")) + assert ct.ensure_cloudflared() == str(cached) + assert cached.read_bytes() == b"MZ" + + +def test_ensure_macos_extracts_tgz_and_chmods(monkeypatch, tmp_path): + cached = tmp_path / "cloudflared" + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-darwin-arm64.tgz", True)) + monkeypatch.setattr(ct, "_cache_path", lambda: cached) + monkeypatch.setattr(ct.sys, "platform", "darwin") + + def fake_download(url, dest): + # dest is cached.with_suffix(".tgz"); write a real archive there. + assert url.endswith("/cloudflared-darwin-arm64.tgz") + with tarfile.open(dest, "w:gz") as tar: + data = b"mach-o" + info = tarfile.TarInfo(name = "cloudflared") + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + return True + + monkeypatch.setattr(ct, "_download", fake_download) + path = ct.ensure_cloudflared() + assert path == str(cached) + assert cached.read_bytes() == b"mach-o" + assert cached.stat().st_mode & 0o111 # chmod applied on posix + assert not cached.with_suffix(".tgz").exists() # temp archive cleaned up + + +# ── .tgz extraction (darwin) ───────────────────────────────────────── + + +def _make_tgz( + tmp_path, + member_name, + data = b"bin", +): + tgz = tmp_path / "cf.tgz" + with tarfile.open(tgz, "w:gz") as tar: + info = tarfile.TarInfo(name = member_name) + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + return tgz + + +def test_tgz_extraction_extracts_clean_member(tmp_path): + tgz = _make_tgz(tmp_path, "cloudflared") + dest = tmp_path / "out" + assert ct._extract_tgz_member(tgz, dest) is True + assert dest.read_bytes() == b"bin" + + +def test_tgz_extraction_rejects_traversal(tmp_path): + tgz = _make_tgz(tmp_path, "../cloudflared") + dest = tmp_path / "out" + assert ct._extract_tgz_member(tgz, dest) is False + assert not dest.exists() + + +def test_tgz_extraction_missing_member(tmp_path): + tgz = _make_tgz(tmp_path, "README") + dest = tmp_path / "out" + assert ct._extract_tgz_member(tgz, dest) is False + + +# ── tunnel lifecycle ───────────────────────────────────────────────── + + +class _FakePopen: + def __init__(self): + self.terminated = False + self.killed = False + self._alive = True + + def poll(self): + return None if self._alive else 0 + + def terminate(self): + self.terminated = True + self._alive = False + + def wait(self, timeout = None): + if self._alive: + raise ct.subprocess.TimeoutExpired(cmd = "cloudflared", timeout = timeout) + return 0 + + def kill(self): + self.killed = True + self._alive = False + + +def test_stop_terminates_process(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + fake = _FakePopen() + t._proc = fake + t.stop() + assert fake.terminated is True + assert t._proc is None + # second stop is a no-op (idempotent) + t.stop() + + +def test_wait_for_url_times_out_without_blocking(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + assert t.wait_for_url(timeout = 0.05) is None + + +def test_start_studio_tunnel_no_binary(monkeypatch): + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: None) + assert ct.start_studio_tunnel(8080) is None + + +def test_start_studio_tunnel_registers_before_wait(monkeypatch): + # The tunnel must be visible to stop_studio_tunnel() during the URL wait, + # else a shutdown in that window orphans cloudflared. + seen = {} + + class _Stub: + def __init__(self, port, binary): + self.url = None + + def start(self): + pass + + def wait_for_url(self, timeout): + seen["active_during_wait"] = ct._active_tunnel is self + self.url = "https://x.trycloudflare.com" + return self.url + + def stop(self): + seen["stopped"] = True + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + try: + assert ct.start_studio_tunnel(8080) == "https://x.trycloudflare.com" + assert seen["active_during_wait"] is True + finally: + ct.stop_studio_tunnel() + + +def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch): + seen = {} + + class _Stub: + def __init__(self, port, binary): + self.url = None + + def start(self): + pass + + def wait_for_url(self, timeout): + return None + + def stop(self): + seen["stopped"] = True + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert seen.get("stopped") is True + assert ct._active_tunnel is None + + +def test_start_studio_tunnel_returns_url(monkeypatch): + class _StubTunnel: + def __init__(self, port, binary): + self.url = None + + def start(self): + self.url = "https://stub-xyz.trycloudflare.com" + + def wait_for_url(self, timeout): + return self.url + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _StubTunnel) + try: + assert ct.start_studio_tunnel(8080) == "https://stub-xyz.trycloudflare.com" + finally: + ct.stop_studio_tunnel() + + +# ── run.py source-level pins (AST / source, no heavy import) ───────── + + +def _func_param_defaults(source, func_name): + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: + args = node.args.args + defaults = node.args.defaults + offset = len(args) - len(defaults) + out = {} + for i, d in enumerate(defaults): + if isinstance(d, ast.Constant): + out[args[offset + i].arg] = d.value + return out + return {} + + +def _argparse_default(source, option): + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if node.func.attr == "add_argument" and node.args: + a0 = node.args[0] + if isinstance(a0, ast.Constant) and a0.value == option: + for kw in node.keywords: + if kw.arg == "default" and isinstance(kw.value, ast.Constant): + return kw.value.value + return None + + +def test_run_server_cloudflare_default_true(): + defaults = _func_param_defaults(_RUN_PY.read_text(), "run_server") + assert defaults.get("cloudflare") is True + + +def test_argparse_cloudflare_default_true(): + assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True + + +def test_run_server_gates_tunnel_on_wildcard(): + # Guard against accidentally widening the trigger beyond 0.0.0.0. + source = _RUN_PY.read_text() + assert "_cloudflare_enabled" in source + assert 'host == "0.0.0.0"' in source + + +def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable): + """Exec the real _print_cloudflare_line source in isolation (run.py has heavy + deps), with the two module globals injected and startup_banner stubbed.""" + src = _RUN_PY.read_text() + tree = ast.parse(src) + func_src = next( + ast.get_source_segment(src, n) + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_print_cloudflare_line" + ) + stub = types.ModuleType("startup_banner") + stub.stdout_supports_color = lambda: False + monkeypatch.setitem(sys.modules, "startup_banner", stub) + captured: list[str] = [] + ns = { + "_cloudflare_url": cloudflare_url, + "_public_reachable": public_reachable, + "print": lambda *a, **k: captured.append(" ".join(str(x) for x in a)), + } + exec(compile(func_src, "", "exec"), ns) + ns["_print_cloudflare_line"]() + return "\n".join(captured) + + +def test_cloudflare_line_reworded_when_public_unreachable(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = False + ) + assert "Use the secure link access via Cloudflare instead: https://x.trycloudflare.com" in out + + +def test_cloudflare_line_default_wording_when_reachable(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = True + ) + assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out + assert "Use the secure link" not in out + + +def test_cloudflare_line_default_wording_when_unknown(monkeypatch): + # Probe did not run / could not decide -> keep the existing wording. + out = _run_print_cloudflare_line( + monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = None + ) + assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out + assert "Use the secure link" not in out + + +def test_cloudflare_line_prints_nothing_without_tunnel(monkeypatch): + out = _run_print_cloudflare_line(monkeypatch, cloudflare_url = None, public_reachable = False) + assert out == "" diff --git a/studio/backend/tests/test_context_overflow_truncation.py b/studio/backend/tests/test_context_overflow_truncation.py new file mode 100644 index 0000000000..4f4c240934 --- /dev/null +++ b/studio/backend/tests/test_context_overflow_truncation.py @@ -0,0 +1,277 @@ +# 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 the opt-in ``context_overflow="truncate_middle"`` passthrough policy. + +On ``exceed_context_size_error`` the passthrough drops middle turn-groups and +retries inside the real window instead of surfacing a fatal 400. Truncation +keeps the system prompt, the first turn, and recent turns, and never orphans +a tool result from its tool_calls turn. Also covers ``/v1/models`` exposing +the real post-readback context window. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from routes.inference import ( + _apply_overflow_truncation, + _clip_long_contents, + _CLIP_MARKER, + _estimate_message_tokens, + _openai_model_objects, + _overflow_truncation_requested, + _parse_overflow_counts, + _truncate_middle_messages, +) +import routes.inference as routes_mod + + +# Nick's actual error body from the Discord report logs. +_NICK_ERROR = ( + '{"detail":"llama-server error: {\\"error\\":{\\"code\\":400,' + '\\"message\\":\\"request (70494 tokens) exceeds the available context size ' + '(67584 tokens), try increasing it\\",\\"type\\":\\"exceed_context_size_error\\",' + '\\"n_prompt_tokens\\":70494,\\"n_ctx\\":67584}}"}' +) + + +def _tool_turn(i: int, result_chars: int = 400) -> list[dict]: + """An assistant tool_calls turn paired with its tool result.""" + return [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": f"call_{i}", + "type": "function", + "function": {"name": "read", "arguments": f'{{"filePath":"/f{i}"}}'}, + } + ], + }, + {"role": "tool", "tool_call_id": f"call_{i}", "content": "x" * result_chars}, + ] + + +def _conversation(n_tool_turns: int = 12) -> list[dict]: + msgs = [ + {"role": "system", "content": "You are an agent." * 20}, + {"role": "user", "content": "Do the big task." * 20}, + ] + for i in range(n_tool_turns): + msgs.extend(_tool_turn(i)) + msgs.append({"role": "assistant", "content": "halfway summary"}) + msgs.append({"role": "user", "content": "keep going"}) + return msgs + + +# --------------------------------------------------------------------------- +# _parse_overflow_counts +# --------------------------------------------------------------------------- + + +def test_parse_overflow_counts_nick_error(): + assert _parse_overflow_counts(_NICK_ERROR) == (70494, 67584) + + +def test_parse_overflow_counts_missing_fields(): + assert _parse_overflow_counts('{"error":"something else"}') is None + + +# --------------------------------------------------------------------------- +# _truncate_middle_messages +# --------------------------------------------------------------------------- + + +def test_truncation_drops_middle_keeps_anchors(): + msgs = _conversation() + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.5) + assert dropped > 0 + assert len(new) == len(msgs) - dropped + # System prompt and task anchor survive. + assert new[0]["role"] == "system" + assert new[1] == msgs[1] + # The most recent turns survive verbatim. + assert new[-1] == msgs[-1] + assert new[-2] == msgs[-2] + + +def test_truncation_never_orphans_tool_results(): + msgs = _conversation() + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.4) + assert dropped > 0 + surviving_call_ids = { + tc["id"] for m in new if m.get("role") == "assistant" for tc in (m.get("tool_calls") or []) + } + for m in new: + if m.get("role") == "tool": + assert m["tool_call_id"] in surviving_call_ids + + +def test_truncation_reduces_estimated_size_toward_target(): + msgs = _conversation() + total = sum(_estimate_message_tokens(m) for m in msgs) + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.5) + new_total = sum(_estimate_message_tokens(m) for m in new) + assert dropped > 0 + assert new_total < total + # Should land at or below the requested share, modulo one whole group. + biggest_group = max( + _estimate_message_tokens(a) + _estimate_message_tokens(b) + for a, b in zip(msgs[2:-2:2], msgs[3:-2:2]) + ) + assert new_total <= int(total * 0.5) + biggest_group + + +def test_truncation_noop_when_keep_ratio_full(): + msgs = _conversation() + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 1.0) + assert dropped == 0 + assert new == msgs + + +def test_truncation_noop_when_only_protected_turns_remain(): + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + *_tool_turn(0), + {"role": "user", "content": "latest"}, + ] + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.1) + assert dropped == 0 + assert new == msgs + + +# --------------------------------------------------------------------------- +# _apply_overflow_truncation +# --------------------------------------------------------------------------- + + +def test_apply_overflow_truncation_mutates_body_and_clamps_max_tokens(): + body = {"messages": _conversation(), "max_tokens": 32000} + assert _apply_overflow_truncation(body, _NICK_ERROR) is True + assert len(body["messages"]) < len(_conversation()) + # Generation headroom: max_tokens clamped to the non-prompt share of n_ctx. + assert body["max_tokens"] <= max(1024, int(67584 * 0.25)) + + +def test_apply_overflow_truncation_returns_false_when_nothing_droppable(): + body = { + "messages": [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + {"role": "user", "content": "latest"}, + ], + "max_tokens": 32000, + } + assert _apply_overflow_truncation(body, _NICK_ERROR) is False + + +def test_apply_overflow_truncation_clips_giant_protected_tool_results(): + """One giant burst (few turn-groups, all protected) must still shrink: + stage 2 clips oversized tool contents instead of giving up.""" + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + *_tool_turn(0, result_chars = 60000), + *_tool_turn(1, result_chars = 60000), + ] + body = {"messages": msgs, "max_tokens": 32000} + n_before = len(msgs) + assert _apply_overflow_truncation(body, _NICK_ERROR) is True + # No message disappeared (pairing intact), but contents were clipped. + assert len(body["messages"]) == n_before + clipped = [m for m in body["messages"] if _CLIP_MARKER in str(m.get("content"))] + assert clipped, "expected at least one clipped tool result" + surviving_call_ids = { + tc["id"] + for m in body["messages"] + if m.get("role") == "assistant" + for tc in (m.get("tool_calls") or []) + } + for m in body["messages"]: + if m.get("role") == "tool": + assert m["tool_call_id"] in surviving_call_ids + + +def test_clip_long_contents_reaches_target_and_keeps_structure(): + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + *_tool_turn(0, result_chars = 40000), + {"role": "user", "content": "latest question"}, + ] + total = sum(_estimate_message_tokens(m) for m in msgs) + clipped = _clip_long_contents(msgs, target_est = total // 4) + assert clipped >= 1 + assert sum(_estimate_message_tokens(m) for m in msgs) <= total // 4 + # Roles and count unchanged; the short final user message untouched. + assert [m["role"] for m in msgs] == ["system", "user", "assistant", "tool", "user"] + assert msgs[-1]["content"] == "latest question" + + +def test_overflow_truncation_requested_reads_field(monkeypatch): + monkeypatch.delenv("UNSLOTH_CONTEXT_OVERFLOW", raising = False) + + class _P: + context_overflow = "truncate_middle" + + class _Q: + context_overflow = None + + assert _overflow_truncation_requested(_P()) is True + assert _overflow_truncation_requested(_Q()) is False + assert _overflow_truncation_requested(object()) is False + + +def test_overflow_truncation_server_default_env(monkeypatch): + """UNSLOTH_CONTEXT_OVERFLOW enables the policy for clients that cannot + send custom body fields; an explicit per-request 'error' still wins.""" + + class _Unset: + context_overflow = None + + class _ExplicitError: + context_overflow = "error" + + monkeypatch.setenv("UNSLOTH_CONTEXT_OVERFLOW", "truncate_middle") + assert _overflow_truncation_requested(_Unset()) is True + assert _overflow_truncation_requested(_ExplicitError()) is False + monkeypatch.setenv("UNSLOTH_CONTEXT_OVERFLOW", "error") + assert _overflow_truncation_requested(_Unset()) is False + + +# --------------------------------------------------------------------------- +# /v1/models context metadata +# --------------------------------------------------------------------------- + + +class _FakeLlamaBackend: + is_loaded = True + model_identifier = "unsloth/Qwen3.6-27B-GGUF" + context_length = 67584 + max_context_length = 262144 + + +class _FakeEmptyBackend: + active_model_name = None + + +def test_v1_models_exposes_real_context_window(monkeypatch): + monkeypatch.setattr(routes_mod, "get_llama_cpp_backend", lambda: _FakeLlamaBackend()) + monkeypatch.setattr(routes_mod, "get_inference_backend", lambda: _FakeEmptyBackend()) + models = _openai_model_objects() + assert len(models) == 1 + entry = models[0] + assert entry["id"] == "unsloth/Qwen3.6-27B-GGUF" + # The REAL (post /props readback) window, not the requested one. + assert entry["context_length"] == 67584 + assert entry["max_context_length"] == 262144 diff --git a/studio/backend/tests/test_cpu_threads.py b/studio/backend/tests/test_cpu_threads.py new file mode 100644 index 0000000000..2930c9f081 --- /dev/null +++ b/studio/backend/tests/test_cpu_threads.py @@ -0,0 +1,152 @@ +# 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 Studio's early CPU thread-pool configuration.""" + +import ast +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from utils.cpu_threads import _THREAD_POOL_ENV_VARS, configure_cpu_threads + + +_BACKEND_DIR = Path(__file__).resolve().parent.parent +_RUN_PY = _BACKEND_DIR / "run.py" +_MAIN_PY = _BACKEND_DIR / "main.py" + + +# Explicit positive integers seed all four native pool env vars. +def test_cpu_thread_cap_seeds_native_pool_limits(): + env = {"UNSLOTH_CPU_THREADS": " 6 "} + + configure_cpu_threads(env) + + assert {variable: env[variable] for variable in _THREAD_POOL_ENV_VARS} == { + variable: "6" for variable in _THREAD_POOL_ENV_VARS + } + + +# Explicit per-library values win over the Studio knob via setdefault. +def test_cpu_thread_cap_preserves_runtime_specific_override(): + env = {"UNSLOTH_CPU_THREADS": "4", "OMP_NUM_THREADS": "2"} + + configure_cpu_threads(env) + + assert env["OMP_NUM_THREADS"] == "2" + assert env["MKL_NUM_THREADS"] == "4" + + +# Whitespace / plus-prefix / leading zero all normalise via int(). +@pytest.mark.parametrize("raw", ["+4", "007", " 4 "]) +def test_cpu_thread_cap_normalises_valid_inputs(raw): + env = {"UNSLOTH_CPU_THREADS": raw} + + configure_cpu_threads(env) + + assert env["OMP_NUM_THREADS"] == str(int(raw.strip())) + + +# Unset / empty / whitespace -> no env mutation (pure opt-in). +@pytest.mark.parametrize("raw", [None, "", " ", "\t"]) +def test_cpu_thread_cap_is_opt_in(raw): + env = {} if raw is None else {"UNSLOTH_CPU_THREADS": raw} + snapshot = dict(env) + + configure_cpu_threads(env) + + assert env == snapshot + assert all(variable not in env for variable in _THREAD_POOL_ENV_VARS) + + +# Anything that is not a positive integer raises a clear ValueError. +@pytest.mark.parametrize("raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]) +def test_cpu_thread_cap_requires_positive_integer(raw): + with pytest.raises(ValueError, match = "must be a positive integer"): + configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw}) + + +# env=None path uses real os.environ (production call from run.py / main.py). +def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch): + for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"): + monkeypatch.delenv(variable, raising = False) + monkeypatch.setenv("UNSLOTH_CPU_THREADS", "3") + + configure_cpu_threads() + + for variable in _THREAD_POOL_ENV_VARS: + assert os.environ[variable] == "3" + + +# Calling twice must not flip any seeded value. +def test_cpu_thread_cap_idempotent(monkeypatch): + for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"): + monkeypatch.delenv(variable, raising = False) + monkeypatch.setenv("UNSLOTH_CPU_THREADS", "5") + + configure_cpu_threads() + snapshot = {v: os.environ.get(v) for v in _THREAD_POOL_ENV_VARS} + configure_cpu_threads() + + assert {v: os.environ.get(v) for v in _THREAD_POOL_ENV_VARS} == snapshot + + +def _ast_line_of_configure_call(source: str) -> int: + tree = ast.parse(source) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "configure_cpu_threads" + ): + return node.lineno + raise AssertionError("configure_cpu_threads() call not found") + + +def _ast_line_of_platform_compat_import(source: str) -> int: + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "_platform_compat": + return node.lineno + raise AssertionError("_platform_compat import not found") + + +# AST ordering: configure_cpu_threads() must precede _platform_compat in both +# run.py and main.py. Robust to formatting / line shifts. +@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY]) +def test_cpu_thread_configuration_runs_before_backend_imports(entry_point): + source = entry_point.read_text() + call_line = _ast_line_of_configure_call(source) + compat_line = _ast_line_of_platform_compat_import(source) + assert call_line < compat_line, ( + f"{entry_point.name}: configure_cpu_threads() (line {call_line}) " + f"must precede import _platform_compat (line {compat_line})" + ) + + +# Invalid env -> exit 1, one-line stderr, no traceback, gated before any +# heavy import. Parametrised over both entry points. +@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY]) +def test_invalid_cpu_thread_cap_exits_without_traceback(entry_point): + env = os.environ.copy() + env["UNSLOTH_CPU_THREADS"] = "not-a-count" + + result = subprocess.run( + [sys.executable, str(entry_point)], + env = env, + capture_output = True, + text = True, + ) + + assert result.returncode == 1 + assert ( + "Error: Invalid UNSLOTH_CPU_THREADS value 'not-a-count': " + "UNSLOTH_CPU_THREADS must be a positive integer" + ) in result.stderr + assert "Traceback" not in result.stderr + assert "_platform_compat" not in result.stderr diff --git a/studio/backend/tests/test_dataset_upload_limits.py b/studio/backend/tests/test_dataset_upload_limits.py new file mode 100644 index 0000000000..dc6030edbf --- /dev/null +++ b/studio/backend/tests/test_dataset_upload_limits.py @@ -0,0 +1,63 @@ +# 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 training dataset upload limits and cleanup.""" + +import asyncio +import sys +from pathlib import Path +from typing import cast + +import pytest +from fastapi import HTTPException, UploadFile + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from routes import datasets as datasets_route # noqa: E402 + + +class FakeUploadFile: + def __init__(self, filename: str, chunks: list[bytes]): + self.filename = filename + self._chunks = list(chunks) + + async def read(self, _size: int = -1) -> bytes: + if not self._chunks: + return b"" + return self._chunks.pop(0) + + +@pytest.fixture(autouse = True) +def isolate_upload_dir(tmp_path, monkeypatch): + monkeypatch.setattr(datasets_route, "DATASET_UPLOAD_DIR", tmp_path) + monkeypatch.setattr(datasets_route, "get_upload_limit_bytes", lambda: 1024 * 1024) + monkeypatch.setattr(datasets_route, "get_upload_limit_label", lambda: "1MB") + return tmp_path + + +def test_dataset_upload_under_configured_cap_succeeds(isolate_upload_dir): + upload = FakeUploadFile("sample.csv", [b"a,b\n1,2\n"]) + response = asyncio.run( + datasets_route.upload_dataset(cast(UploadFile, upload), current_subject = "test-user") + ) + stored = Path(response.stored_path) + assert response.filename == "sample.csv" + assert stored.exists() + assert stored.parent == isolate_upload_dir + assert stored.read_bytes() == b"a,b\n1,2\n" + + +def test_dataset_upload_over_configured_cap_removes_partial_file(isolate_upload_dir): + upload = FakeUploadFile( + "sample.csv", + [b"x" * (1024 * 1024), b"y"], + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + datasets_route.upload_dataset(cast(UploadFile, upload), current_subject = "test-user") + ) + assert exc.value.status_code == 413 + assert "Maximum is 1MB" in exc.value.detail + assert list(isolate_upload_dir.iterdir()) == [] diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index bc8788fd90..b4b39ba1a9 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -9,7 +9,7 @@ import sqlite3 import subprocess import sys from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import jwt import pytest @@ -51,12 +51,8 @@ def auth_client(): def data_recipe_jobs_module(): - route_path = ( - Path(__file__).resolve().parents[1] / "routes" / "data_recipe" / "jobs.py" - ) - spec = importlib.util.spec_from_file_location( - "_desktop_data_recipe_jobs", route_path - ) + route_path = Path(__file__).resolve().parents[1] / "routes" / "data_recipe" / "jobs.py" + spec = importlib.util.spec_from_file_location("_desktop_data_recipe_jobs", route_path) jobs_route = importlib.util.module_from_spec(spec) assert spec.loader is not None spec.loader.exec_module(jobs_route) @@ -243,7 +239,7 @@ def test_consume_refresh_token_second_call_returns_none(): def test_consume_refresh_token_concurrent_only_one_succeeds(tmp_path, monkeypatch): - """64-thread pile-up against one token; DELETE RETURNING permits one winner.""" + """64-thread pile-up on one token; DELETE RETURNING permits one winner.""" seed_user() from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone @@ -258,16 +254,14 @@ def test_consume_refresh_token_concurrent_only_one_succeeds(tmp_path, monkeypatc try: return storage.consume_refresh_token(raw) except sqlite3.OperationalError: - # "database is locked" under heavy contention; treat as losing the race. + # "database is locked" under contention; treat as losing the race. return None with ThreadPoolExecutor(max_workers = workers) as pool: results = list(pool.map(attempt, range(workers))) successes = [r for r in results if r is not None] - assert ( - len(successes) == 1 - ), f"expected exactly one consumer to win, got {len(successes)}" + assert len(successes) == 1, f"expected exactly one consumer to win, got {len(successes)}" assert successes[0] == (storage.DEFAULT_ADMIN_USERNAME, False) @@ -285,9 +279,7 @@ def test_desktop_session_uses_real_admin_identity_for_api_keys(): seed_user(must_change_password = True) raw = storage.create_desktop_secret() client = auth_client() - token = client.post("/api/auth/desktop-login", json = {"secret": raw}).json()[ - "access_token" - ] + token = client.post("/api/auth/desktop-login", json = {"secret": raw}).json()["access_token"] response = client.post( "/api/auth/api-keys", @@ -302,8 +294,7 @@ def test_desktop_session_uses_real_admin_identity_for_api_keys(): def test_local_recipe_token_authenticates_as_admin_for_desktop_user(loaded_local_model): # _inject_local_providers mints an internal sk-unsloth-* API key (not a - # forwarded JWT). The unified API-key path validates as the real admin - # user regardless of whether the incoming session was desktop or web. + # forwarded JWT) that validates as admin whether the session was desktop or web. from auth.authentication import create_access_token, get_current_subject seed_user(must_change_password = True) @@ -322,14 +313,11 @@ def test_local_recipe_token_authenticates_as_admin_for_desktop_user(loaded_local scheme = "Bearer", credentials = local_token, ) - assert ( - asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME - ) + assert asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_model): - # Mirror of the desktop variant: API-key issuance is identical for web - # and desktop incoming tokens; auth via get_current_subject works the same. + # Mirror of the desktop variant: API-key issuance is identical for web/desktop tokens. from auth.authentication import create_access_token, get_current_subject seed_user(must_change_password = False) @@ -345,9 +333,7 @@ def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_mod scheme = "Bearer", credentials = local_token, ) - assert ( - asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME - ) + assert asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME def test_desktop_login_rejects_invalid_secret(): @@ -429,20 +415,39 @@ def test_desktop_capabilities_json_reports_rollout_safe_flags(): def test_health_response_reports_desktop_capability_fields(monkeypatch): - router_stub = SimpleNamespace( - auth_router = APIRouter(), - chat_history_router = APIRouter(), - data_recipe_router = APIRouter(), - datasets_router = APIRouter(), - export_router = APIRouter(), - inference_router = APIRouter(), - inference_studio_router = APIRouter(), - models_router = APIRouter(), - providers_router = APIRouter(), - training_history_router = APIRouter(), - training_router = APIRouter(), - ) - monkeypatch.setitem(sys.modules, "routes", router_stub) + routes_module = ModuleType("routes") + routes_module.__path__ = [] + settings_module = ModuleType("routes.settings") + settings_module.router = APIRouter() + llama_module = ModuleType("routes.llama") + llama_module.router = APIRouter() + prompts_module = ModuleType("routes.prompts") + prompts_module.router = APIRouter() + + for name, router in { + "auth_router": APIRouter(), + "chat_history_router": APIRouter(), + "data_recipe_router": APIRouter(), + "datasets_router": APIRouter(), + "export_router": APIRouter(), + "inference_router": APIRouter(), + "inference_studio_router": APIRouter(), + "mcp_servers_router": APIRouter(), + "models_router": APIRouter(), + "providers_router": APIRouter(), + "rag_router": APIRouter(), + "settings_router": settings_module.router, + "training_history_router": APIRouter(), + "training_router": APIRouter(), + }.items(): + setattr(routes_module, name, router) + routes_module.settings = settings_module + routes_module.llama = llama_module + + monkeypatch.setitem(sys.modules, "routes", routes_module) + monkeypatch.setitem(sys.modules, "routes.settings", settings_module) + monkeypatch.setitem(sys.modules, "routes.llama", llama_module) + monkeypatch.setitem(sys.modules, "routes.prompts", prompts_module) import studio.backend.main as backend_main @@ -469,8 +474,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): def test_provision_desktop_auth_writes_secret_and_creates_db_without_backend_deps( - tmp_path, - monkeypatch, + tmp_path, monkeypatch ): auth_dir = tmp_path / "auth" auth_dir.mkdir() @@ -525,12 +529,9 @@ if result.exit_code != 0: """ ).fetchone() app_secrets = { - row["key"]: row["value"] - for row in conn.execute("SELECT key, value FROM app_secrets") - } - refresh_columns = { - row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)") + row["key"]: row["value"] for row in conn.execute("SELECT key, value FROM app_secrets") } + refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")} finally: conn.close() @@ -622,9 +623,7 @@ def test_update_password_clears_desktop_secret(): raw = storage.create_desktop_secret() assert storage.validate_desktop_secret(raw) == storage.DEFAULT_ADMIN_USERNAME - changed = storage.update_password( - storage.DEFAULT_ADMIN_USERNAME, "new-admin-password" - ) + changed = storage.update_password(storage.DEFAULT_ADMIN_USERNAME, "new-admin-password") assert changed is True assert storage.validate_desktop_secret(raw) is None @@ -640,11 +639,7 @@ def test_update_password_on_unknown_user_leaves_desktop_secret_intact(): def test_desktop_auth_provision_has_bounded_timeout(): rs_path = ( - Path(__file__).resolve().parents[3] - / "studio" - / "src-tauri" - / "src" - / "desktop_auth.rs" + Path(__file__).resolve().parents[3] / "studio" / "src-tauri" / "src" / "desktop_auth.rs" ) src = rs_path.read_text() start = src.index("async fn provision_desktop_auth(") diff --git a/studio/backend/tests/test_detect_mmproj_file.py b/studio/backend/tests/test_detect_mmproj_file.py index cdb73448be..64dd8ebd90 100644 --- a/studio/backend/tests/test_detect_mmproj_file.py +++ b/studio/backend/tests/test_detect_mmproj_file.py @@ -46,7 +46,7 @@ def test_returns_none_when_no_mmproj(tmp_path: Path): def test_single_matching_family_mmproj_picked(tmp_path: Path): - """Single same-family projector: returned (historical behaviour).""" + """Single same-family projector is 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()) @@ -132,10 +132,7 @@ def test_family_token_mistral_does_not_match_ministral(): 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" - ) + assert _detect_family_token("Devstral-Small-2-24B-Instruct-2512-BF16.gguf") == "devstral" def test_family_token_picks_leftmost_when_multiple_present(): @@ -160,9 +157,7 @@ def test_family_token_new_families_recognised(): 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" - ) + 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 diff --git a/studio/backend/tests/test_export_log_cursor.py b/studio/backend/tests/test_export_log_cursor.py index 734ca522c9..5b614ab9c8 100644 --- a/studio/backend/tests/test_export_log_cursor.py +++ b/studio/backend/tests/test_export_log_cursor.py @@ -4,24 +4,15 @@ """ Regression tests for the export log ring-buffer cursor semantics. -Context: the live export log SSE stream has a race where the frontend -opens the SSE connection AFTER the POST that starts the export. Any -lines the worker subprocess emits during the gap between POST and SSE -connect get buffered with seqs 1..k, and then the SSE default cursor -`get_current_log_seq()` returns k -- so lines 1..k are forever -unreachable to that client. +Race: the frontend opens the SSE connection AFTER the POST that starts the +export, so lines emitted in the gap (seqs 1..k) are unreachable when the SSE +default cursor `get_current_log_seq()` returns k. -Fix: `clear_logs()` snapshots the pre-run seq into `_run_start_seq` -(exposed via `get_run_start_seq()`), and `routes/export.py` defaults -the SSE cursor to that snapshot instead of the current seq. Every line -appended during the current run has seq strictly greater than the -snapshot, so the client sees the full run regardless of when it -connects. +Fix: `clear_logs()` snapshots the pre-run seq into `_run_start_seq` (via +`get_run_start_seq()`), and the SSE cursor defaults to that snapshot, so the +client sees the full run regardless of connect time. -These tests exercise the orchestrator-side contract only (no -subprocess, no FastAPI, no frontend). The routes-level integration -with get_run_start_seq() is a one-line edit covered by manual testing -and the frontend build. +These tests exercise the orchestrator contract only (no subprocess/FastAPI). """ from __future__ import annotations @@ -33,27 +24,20 @@ from pathlib import Path import pytest -# Backend root on sys.path so `from core.export.orchestrator import ...` -# and friends resolve without the studio app bootstrap. +# Backend root on sys.path so core.export.orchestrator resolves without bootstrap. _BACKEND_DIR = Path(__file__).resolve().parent.parent if str(_BACKEND_DIR) not in sys.path: sys.path.insert(0, str(_BACKEND_DIR)) -# ExportOrchestrator imports structlog and a few heavy modules at the -# top of orchestrator.py. Stub the ones we don't need in these unit -# tests so the import succeeds on machines without the full studio -# venv. +# Stub orchestrator.py's heavy top-level imports so it loads without the venv. _loggers_stub = types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) -# structlog is only used for a module-level import; a bare stub is -# enough because we never call into it in these tests. +# structlog is only a module-level import here; a bare stub suffices. sys.modules.setdefault("structlog", types.ModuleType("structlog")) -# utils.paths.outputs_root is only called inside scan_checkpoints which -# we don't hit in these tests. Provide a stub module so the top-level -# import in orchestrator.py resolves. +# Stub utils.paths so orchestrator.py's top-level import resolves. _utils_pkg = types.ModuleType("utils") _utils_pkg.__path__ = [] # mark as package _utils_paths_stub = types.ModuleType("utils.paths") @@ -66,11 +50,14 @@ sys.modules.setdefault("utils.paths", _utils_paths_stub) def orchestrator(): """Fresh ExportOrchestrator with only the log-buffer state exercised.""" from core.export.orchestrator import ExportOrchestrator - return ExportOrchestrator() -def _append(orch, line: str, stream: str = "stdout") -> None: +def _append( + orch, + line: str, + stream: str = "stdout", +) -> None: """Shortcut for simulating a worker log message.""" orch._append_log({"type": "log", "stream": stream, "line": line, "ts": 0.0}) @@ -81,14 +68,14 @@ def _append(orch, line: str, stream: str = "stdout") -> None: def test_run_start_seq_is_zero_before_any_logs(orchestrator) -> None: - """A brand-new orchestrator must report run_start_seq == 0 so a - first SSE connection picks up every line from seq 1 onward.""" + """run_start_seq == 0 on a new orchestrator, so the first SSE picks up + every line from seq 1 onward.""" assert orchestrator.get_run_start_seq() == 0 def test_clear_logs_snapshots_current_seq(orchestrator) -> None: - """clear_logs() must capture _log_seq BEFORE clearing the buffer, - so subsequent runs can anchor their SSE cursor at the snapshot.""" + """clear_logs() captures _log_seq BEFORE clearing, so the next run can + anchor its SSE cursor at the snapshot.""" _append(orchestrator, "old run line 1") _append(orchestrator, "old run line 2") _append(orchestrator, "old run line 3") @@ -106,14 +93,9 @@ def test_clear_logs_snapshots_current_seq(orchestrator) -> None: def test_sse_default_cursor_catches_all_current_run_lines(orchestrator) -> None: - """Simulate the POST-then-SSE race: worker starts emitting lines - immediately after clear_logs(), SSE connects several lines later. - Using get_run_start_seq() as the default cursor MUST return every - line emitted since clear_logs() ran. - - Pre-fix, the SSE defaulted to get_current_log_seq() at connect - time, which would return the last-seen seq and miss lines N+1..M. - """ + """POST-then-SSE race: worker emits lines right after clear_logs(), SSE + connects later. get_run_start_seq() as the default cursor returns every + line since clear_logs() (pre-fix it missed them).""" # Previous run leaves some buffered lines. _append(orchestrator, "previous run line A") _append(orchestrator, "previous run line B") @@ -127,8 +109,7 @@ def test_sse_default_cursor_catches_all_current_run_lines(orchestrator) -> None: _append(orchestrator, "Loading checkpoint: /foo/bar") _append(orchestrator, "Starting export...") - # SSE connects now and asks "give me everything after the run - # start cursor". + # SSE connects and asks for everything after the run start cursor. entries, new_cursor = orchestrator.get_logs_since(run_start) # All three early lines must be present. Pre-fix this was []. @@ -142,10 +123,8 @@ def test_sse_default_cursor_catches_all_current_run_lines(orchestrator) -> None: def test_sse_default_cursor_excludes_previous_run(orchestrator) -> None: - """After clear_logs(), lines from the PREVIOUS run must not leak - into the new run's SSE stream. Pre-fix this worked correctly - (clear_logs cleared the deque); the fix must preserve it. - """ + """After clear_logs(), previous-run lines must not leak into the new + run's SSE stream (the fix must preserve this).""" _append(orchestrator, "previous run line 1") _append(orchestrator, "previous run line 2") _append(orchestrator, "previous run line 3") @@ -161,10 +140,8 @@ def test_sse_default_cursor_excludes_previous_run(orchestrator) -> None: def test_clear_logs_twice_advances_run_start(orchestrator) -> None: - """Back-to-back clear_logs() calls (e.g. cleanup -> load -> - export in the same dialog session) must each re-anchor run_start - at the current seq, so successive runs each start with a fresh - low-water mark.""" + """Back-to-back clear_logs() calls each re-anchor run_start at the current + seq, giving successive runs a fresh low-water mark.""" _append(orchestrator, "run 1 line a") _append(orchestrator, "run 1 line b") diff --git a/studio/backend/tests/test_external_provider_proxy_env.py b/studio/backend/tests/test_external_provider_proxy_env.py new file mode 100644 index 0000000000..f17b655908 --- /dev/null +++ b/studio/backend/tests/test_external_provider_proxy_env.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from pathlib import Path +import importlib.util +import sys + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_EXTERNAL_PROVIDER_PATH = ( + Path(__file__).resolve().parent.parent / "core/inference/external_provider.py" +) + + +def _load_external_provider_module(): + spec = importlib.util.spec_from_file_location( + "external_provider_under_test", + _EXTERNAL_PROVIDER_PATH, + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_shared_http_client_ignores_unsupported_proxy_scheme(monkeypatch): + ep_mod = _load_external_provider_module() + calls = [] + + class FakeAsyncClient: + def __init__(self, **kwargs): + calls.append(kwargs) + if kwargs.get("trust_env") is not False: + raise ValueError("Unknown scheme for proxy URL URL('socks4://127.0.0.1:12345')") + + monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient) + + client = ep_mod._create_shared_http_client() + + assert isinstance(client, FakeAsyncClient) + assert calls == [{}, {"trust_env": False}] + + +def test_shared_http_client_ignores_missing_socksio(monkeypatch): + ep_mod = _load_external_provider_module() + calls = [] + + class FakeAsyncClient: + def __init__(self, **kwargs): + calls.append(kwargs) + if kwargs.get("trust_env") is not False: + raise ImportError( + "Using SOCKS proxy, but the 'socksio' package is not installed. " + "Make sure to install httpx using `pip install httpx[socks]`." + ) + + monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient) + + client = ep_mod._create_shared_http_client() + + assert isinstance(client, FakeAsyncClient) + assert calls == [{}, {"trust_env": False}] + + +def test_shared_http_client_reraises_other_value_errors(monkeypatch): + ep_mod = _load_external_provider_module() + + class FakeAsyncClient: + def __init__(self, **kwargs): + raise ValueError("different httpx setup error") + + monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient) + + try: + ep_mod._create_shared_http_client() + except ValueError as exc: + assert str(exc) == "different httpx setup error" + else: + raise AssertionError("expected ValueError") diff --git a/studio/backend/tests/test_external_provider_usage_chunk.py b/studio/backend/tests/test_external_provider_usage_chunk.py index 82c641d049..ebfd6a8f50 100644 --- a/studio/backend/tests/test_external_provider_usage_chunk.py +++ b/studio/backend/tests/test_external_provider_usage_chunk.py @@ -1,27 +1,12 @@ # 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 prompt-cache accounting chunk emitted by the external- -provider streaming proxy. +"""Unit tests for the prompt-cache accounting chunk from the external-provider proxy. -The streaming Anthropic + OpenAI Responses paths now emit one extra -``include_usage``-style SSE chunk (``choices: []`` with a populated -``usage`` block) just before ``[DONE]`` / after the final -``finish_reason`` chunk. This lets clients surface cache savings -without scraping the structlog stream. - -Covers: -- Helper alone: shape for Anthropic / OpenAI usage payloads, missing - fields treated as 0, all-zero usage suppressed. -- Anthropic stream: ``message_start.usage`` + ``message_delta.usage`` - with ``cache_creation_input_tokens`` and ``cache_read_input_tokens`` - produce the expected usage chunk before ``[DONE]``. -- OpenAI Responses stream: ``response.completed.usage`` with - ``input_tokens_details.cached_tokens`` produces the expected usage - chunk after the ``stop`` finish_reason chunk. -- OpenAI Responses ``response.incomplete`` also emits the usage chunk - so length-truncated turns still report cached tokens. +The streaming Anthropic + OpenAI Responses paths emit one extra include_usage +SSE chunk (``choices: []`` with a ``usage`` block) before ``[DONE]`` so clients +see cache savings. Covers the helper directly plus the Anthropic stream and the +OpenAI Responses completed/incomplete streams. """ import asyncio @@ -57,9 +42,9 @@ def test_build_usage_chunk_anthropic_shape(): assert payload["object"] == "chat.completion.chunk" assert payload["choices"] == [] usage = payload["usage"] - # Anthropic's input_tokens excludes cache buckets; prompt_tokens - # must add all three input components together so downstream - # context / cost displays see the real prompt size. + # Anthropic's input_tokens excludes cache buckets; prompt_tokens must + # sum all three input components so downstream context/cost displays + # see the real prompt size. assert usage["prompt_tokens"] == 8 + 1367 + 18901 assert usage["completion_tokens"] == 862 assert usage["total_tokens"] == 8 + 1367 + 18901 + 862 @@ -92,9 +77,8 @@ def test_build_usage_chunk_openai_shape(): def test_build_usage_chunk_missing_fields_default_to_zero(): - # OpenAI Responses can return a usage object without - # input_tokens_details when prompt caching is unused; the helper - # should still emit a chunk with cached_tokens=0. + # OpenAI Responses can omit input_tokens_details when prompt caching is + # unused; the helper should still emit a chunk with cached_tokens=0. line = _build_usage_chunk( "chatcmpl-z", "openai", @@ -107,7 +91,7 @@ def test_build_usage_chunk_missing_fields_default_to_zero(): def test_build_usage_chunk_returns_none_when_all_zero(): # If upstream errored before any usage event, suppress the chunk to - # avoid surfacing a misleading "0 tokens" line. + # avoid a misleading "0 tokens" line. assert _build_usage_chunk("id", "anthropic", {}) is None assert _build_usage_chunk("id", "anthropic", None) is None assert _build_usage_chunk("id", "openai", {}) is None @@ -191,11 +175,7 @@ def _usage_chunks(lines: list[str]) -> list[dict]: parsed = json.loads(payload) except json.JSONDecodeError: continue - if ( - isinstance(parsed, dict) - and "usage" in parsed - and parsed.get("choices") == [] - ): + if isinstance(parsed, dict) and "usage" in parsed and parsed.get("choices") == []: out.append(parsed["usage"]) return out @@ -256,13 +236,9 @@ def test_anthropic_stream_emits_usage_chunk_before_done(monkeypatch): # Usage chunk must come before [DONE]. data_lines = [ln for ln in lines if ln.startswith("data:")] - done_idx = next( - i for i, ln in enumerate(data_lines) if ln.strip().endswith("[DONE]") - ) + done_idx = next(i for i, ln in enumerate(data_lines) if ln.strip().endswith("[DONE]")) usage_idx = next( - i - for i, ln in enumerate(data_lines) - if '"usage":' in ln and '"choices": []' in ln + i for i, ln in enumerate(data_lines) if '"usage":' in ln and '"choices": []' in ln ) assert usage_idx < done_idx diff --git a/studio/backend/tests/test_frontend_resolution.py b/studio/backend/tests/test_frontend_resolution.py index 2b49763386..c3e0524a30 100644 --- a/studio/backend/tests/test_frontend_resolution.py +++ b/studio/backend/tests/test_frontend_resolution.py @@ -3,9 +3,8 @@ """Tests for the frontend-dist resolver in studio/backend/run.py. -Loads only the relevant helpers via importlib so the test does not pull in -uvicorn / FastAPI / unsloth's full dependency tree. Pairs with the AST-style -test_host_defaults.py. +Loads only the relevant helpers via importlib to avoid pulling in +uvicorn / FastAPI / unsloth's deps. Pairs with AST-style test_host_defaults.py. """ import ast @@ -19,8 +18,8 @@ _REPO_STUDIO_DIR = _RUN_PY.parent.parent # studio/ def _load_helpers_only(): - """Import just the resolver helpers from run.py without executing the - server-side imports (uvicorn, structlog, etc.).""" + """Import just the resolver helpers from run.py, skipping server-side + imports (uvicorn, structlog, etc.).""" source = _RUN_PY.read_text(encoding = "utf-8") tree = ast.parse(source) keep = [] @@ -91,7 +90,7 @@ def test_resolver_falls_back_to_studio_home_site_packages(tmp_path, monkeypatch) def test_resolver_falls_back_via_editable_pth(tmp_path, monkeypatch): """Simulates a `--local` install: dedicated venv with an editable .pth - pointing at a cloned repo that owns the built dist.""" + pointing at a cloned repo owning the built dist.""" studio_home = tmp_path / "studio_home" sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" sp.mkdir(parents = True) @@ -100,7 +99,7 @@ def test_resolver_falls_back_via_editable_pth(tmp_path, monkeypatch): repo_dist = repo_studio / "frontend" / "dist" repo_dist.mkdir(parents = True) (repo_dist / "index.html").write_text("", encoding = "utf-8") - # Minimal `__editable___pkg_finder.py` carrying a MAPPING dict that + # Minimal `__editable___pkg_finder.py` with the MAPPING dict that # setuptools' editable install generator writes. finder = sp / "__editable___unsloth_0_0_0_finder.py" finder.write_text( @@ -127,16 +126,10 @@ def test_iter_candidates_handles_missing_studio_home(tmp_path, monkeypatch): def test_resolver_falls_back_to_windows_layout_site_packages(tmp_path, monkeypatch): """Pins the `Lib/site-packages` (capital L) Windows venv layout - alongside the POSIX `lib/python*/site-packages` path.""" + alongside the POSIX `lib/python*/site-packages`.""" studio_home = tmp_path / "studio_home" sp_dist = ( - studio_home - / "unsloth_studio" - / "Lib" - / "site-packages" - / "studio" - / "frontend" - / "dist" + studio_home / "unsloth_studio" / "Lib" / "site-packages" / "studio" / "frontend" / "dist" ) sp_dist.mkdir(parents = True) (sp_dist / "index.html").write_text("", encoding = "utf-8") @@ -149,20 +142,19 @@ def test_resolver_falls_back_to_windows_layout_site_packages(tmp_path, monkeypat def test_resolver_does_not_crash_on_non_dict_mapping_literal(tmp_path, monkeypatch): - """A finder file whose MAPPING value is a set / list / non-dict literal - (theoretically possible if the regex matched a brace-delimited literal - that ast.literal_eval can parse) must not AttributeError. The resolver - should skip that finder and keep probing.""" + """A finder whose MAPPING value is a set/list/non-dict literal (possible + if the regex matched a brace-delimited literal ast.literal_eval can parse) + must not AttributeError. The resolver should skip it and keep probing.""" studio_home = tmp_path / "studio_home" sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" sp.mkdir(parents = True) - # Bad finder: set literal, not a dict. ast.literal_eval parses it as set; - # any .get() call on it would raise AttributeError. + # Bad finder: set literal, not a dict. literal_eval parses it as a set, + # so any .get() call on it would raise AttributeError. (sp / "__editable___bad_0_0_0_finder.py").write_text( "MAPPING: dict[str, str] = {'studio', 'unsloth', 'unsloth_cli'}\n", encoding = "utf-8", ) - # Good finder that should still be discovered after the bad one is skipped. + # Good finder, still discovered after the bad one is skipped. repo_root = tmp_path / "clone" repo_dist = repo_root / "studio" / "frontend" / "dist" repo_dist.mkdir(parents = True) @@ -180,9 +172,8 @@ def test_resolver_does_not_crash_on_non_dict_mapping_literal(tmp_path, monkeypat def test_resolver_handles_multiline_mapping_dict(tmp_path, monkeypatch): - """A future setuptools / black reformat that wraps the MAPPING dict - across multiple lines must still parse and resolve. Locks in the - `[^}]*` + re.DOTALL behaviour.""" + """A future setuptools/black reformat wrapping the MAPPING dict across + multiple lines must still parse and resolve. Locks in `[^}]*` + re.DOTALL.""" studio_home = tmp_path / "studio_home" sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" sp.mkdir(parents = True) @@ -210,8 +201,8 @@ def test_resolver_handles_multiline_mapping_dict(tmp_path, monkeypatch): def test_systemexit_message_contains_actionable_fixes(tmp_path, monkeypatch): """The user-facing recovery message is a contract: it must surface the - attempted paths and every concrete fix. Pin its structure so a future - refactor doesn't drop one.""" + attempted paths and every concrete fix. Pin its structure so a refactor + doesn't drop one.""" import os import sys diff --git a/studio/backend/tests/test_gemini_provider.py b/studio/backend/tests/test_gemini_provider.py new file mode 100644 index 0000000000..85ceb04d27 --- /dev/null +++ b/studio/backend/tests/test_gemini_provider.py @@ -0,0 +1,5353 @@ +# 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 native Gemini API translation layer. + +Gemini does NOT speak OpenAI Chat Completions on its primary endpoint +(`streamGenerateContent`). `_stream_gemini` in +`core/inference/external_provider.py` translates between the two shapes: + + Request: + OpenAI messages [{role, content}] + -> Gemini contents [{role, parts: [{text}|{inlineData}|{functionCall}|...]}] + + systemInstruction.parts[].text for role=system messages + + generationConfig.{temperature,topP,topK,maxOutputTokens} + + tools[{googleSearch:{}}] for web_search + + tools[{codeExecution:{}}] for code_execution + + responseModalities=[TEXT,IMAGE] for Nano Banana (gemini-2.5-flash-image) + + cachedContent for prompt caching + + Response: + Gemini SSE chunks { candidates:[{content:{parts:[...]}, finishReason}], + usageMetadata:{promptTokenCount, candidatesTokenCount} } + -> OpenAI chat.completion.chunk frames + (delta.content for text, delta.tool_calls for functionCall, + _toolEvent for image_b64/web_search, usage block before [DONE]) + +These tests pin the outbound body shape AND the inbound translation via +httpx.MockTransport (no live network). Mirrors test_anthropic_cache_ttl.py +and test_openai_image_generation.py. +""" + +import asyncio +import base64 +import json + +import httpx +import pytest + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +_active_mock_clients: list[httpx.AsyncClient] = [] + + +def _drive(coro): + # Fresh loop per drive so tests don't share asyncio state. Close mocked + # clients + shutdown async-generators inside this loop so Python 3.13 + # doesn't emit `Response.aiter_*.aclose was never awaited` on GC. + loop = asyncio.new_event_loop() + try: + result = loop.run_until_complete(coro) + while _active_mock_clients: + mc = _active_mock_clients.pop() + loop.run_until_complete(mc.aclose()) + return result + finally: + try: + loop.run_until_complete(loop.shutdown_asyncgens()) + finally: + loop.close() + + +def _make_gemini_client( + base_url: str = "https://generativelanguage.googleapis.com/v1beta", +) -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "gemini", + base_url = base_url, + api_key = "AIza-test-key", + ) + + +def _mock_http(monkeypatch, handler): + mock_client = httpx.AsyncClient(transport = httpx.MockTransport(handler)) + monkeypatch.setattr(ep_mod, "_http_client", mock_client) + # `_drive` acloses this at end of run inside the same event loop, so we + # don't leak an unawaited aclose() coroutine. + _active_mock_clients.append(mock_client) + + +def _gemini_sse(events: list[dict]) -> bytes: + """Encode a list of dicts as Gemini-style SSE frames (`data:` lines).""" + chunks: list[str] = [] + for event in events: + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def _capture_body(monkeypatch, **kwargs) -> dict: + """Drive a single stream and return the captured outbound request body.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + captured["url"] = str(request.url) + captured["method"] = request.method + # Minimal valid Gemini stream so the helper completes. + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + messages = kwargs.pop("messages", [{"role": "user", "content": "hi"}]) + model = kwargs.pop("model", "gemini-2.5-flash") + temperature = kwargs.pop("temperature", 0.7) + top_p = kwargs.pop("top_p", 0.95) + max_tokens = kwargs.pop("max_tokens", 64) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = messages, + model = model, + temperature = temperature, + top_p = top_p, + max_tokens = max_tokens, + **kwargs, + ): + pass + await client.close() + + _drive(run()) + return captured + + +def _collect(monkeypatch, sse_events, **kwargs) -> list[str]: + """Drive a stream with a custom set of SSE events and return raw lines.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _gemini_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + messages = kwargs.pop("messages", [{"role": "user", "content": "hi"}]) + model = kwargs.pop("model", "gemini-2.5-flash") + temperature = kwargs.pop("temperature", 0.7) + top_p = kwargs.pop("top_p", 0.95) + max_tokens = kwargs.pop("max_tokens", 64) + + out: list[str] = [] + + async def run(): + client = _make_gemini_client() + async for line in client.stream_chat_completion( + messages = messages, + model = model, + temperature = temperature, + top_p = top_p, + max_tokens = max_tokens, + **kwargs, + ): + out.append(line) + await client.close() + + _drive(run()) + return out + + +def _parse_chunks(lines: list[str]) -> list[dict]: + out: list[dict] = [] + for raw in lines: + if not raw.startswith("data:"): + continue + payload = raw[len("data:") :].strip() + if not payload or payload == "[DONE]": + continue + try: + out.append(json.loads(payload)) + except json.JSONDecodeError: + continue + return out + + +# ── request body translation ───────────────────────────────────────── + + +def test_request_body_uses_contents_and_parts_shape(monkeypatch): + """OpenAI messages must be translated to Gemini's `contents` shape.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + {"role": "user", "content": "Follow up"}, + ], + ) + body = captured["body"] + # system -> systemInstruction + assert body["systemInstruction"] == {"parts": [{"text": "Be brief."}]}, body + # user/assistant -> contents with role user/model + assert body["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]}, + {"role": "model", "parts": [{"text": "Hi there"}]}, + {"role": "user", "parts": [{"text": "Follow up"}]}, + ], body["contents"] + # generationConfig fields map across with Google's casing. + gc = body["generationConfig"] + assert gc["temperature"] == 0.7 + assert gc["topP"] == 0.95 + assert gc["maxOutputTokens"] == 64 + + +def test_request_url_targets_stream_generate_content(monkeypatch): + """Helper must POST to /v1beta/models/{model}:streamGenerateContent?alt=sse.""" + captured = _capture_body(monkeypatch, model = "gemini-2.5-pro") + url = captured["url"] + assert ":streamGenerateContent" in url, url + assert "alt=sse" in url, url + assert "/v1beta/models/gemini-2.5-pro" in url, url + assert captured["method"] == "POST" + + +def test_request_auth_header_uses_x_goog_api_key(monkeypatch): + """API key must be sent on `x-goog-api-key`, not Authorization.""" + captured = _capture_body(monkeypatch) + hdrs = captured["headers"] + assert hdrs.get("x-goog-api-key") == "AIza-test-key", hdrs + assert "authorization" not in {k.lower() for k in hdrs}, hdrs + + +def test_top_k_forwarded_only_when_positive(monkeypatch): + """top_k is opt-in; only positive integers reach the wire.""" + captured = _capture_body(monkeypatch, top_k = 40) + assert captured["body"]["generationConfig"]["topK"] == 40 + + captured = _capture_body(monkeypatch, top_k = 0) + assert "topK" not in captured["body"]["generationConfig"] + + +def test_presence_penalty_forwarded_to_generation_config(monkeypatch): + """A non-zero presence_penalty reaches generationConfig.presencePenalty.""" + captured = _capture_body(monkeypatch, presence_penalty = 0.7) + assert captured["body"]["generationConfig"]["presencePenalty"] == 0.7 + + # Default zero is omitted, matching top_k semantics. + captured = _capture_body(monkeypatch, presence_penalty = 0.0) + assert "presencePenalty" not in captured["body"]["generationConfig"] + + +# ── thinkingConfig translation ──────────────────────────────────────── + + +def test_gemini25_flash_thinking_disabled_sets_budget_zero(monkeypatch): + """Gemini 2.5 Flash still uses thinkingBudget; 0 = off.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingBudget": 0}, tc + + +def test_gemini3_flash_thinking_disabled_uses_minimal_level(monkeypatch): + """Gemini 3 Flash uses thinkingLevel; "off" maps to minimal + (Gemini 3 cannot turn thinking fully off).""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "minimal"}, tc + + +def test_gemini25_pro_thinking_disabled_uses_small_budget(monkeypatch): + """Gemini 2.5 Pro 400s on thinkingBudget=0 ("only works in thinking + mode"); coerce to a small positive budget.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-pro", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc is not None and tc.get("thinkingBudget", 0) > 0, tc + + +def test_gemini3_pro_thinking_disabled_uses_low_level(monkeypatch): + """Gemini 3 Pro uses thinkingLevel and rejects 'minimal' (Pro tier), so + 'off' coerces to 'low' (lowest the API accepts).""" + for model in ( + "gemini-3.1-pro-preview", + "gemini-3-pro-preview", + "gemini-3.5-pro", + "gemini-pro-latest", + ): + captured = _capture_body( + monkeypatch, + model = model, + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, (model, tc) + + +def test_gemini25_flash_effort_levels_map_to_budgets(monkeypatch): + """Gemini 2.5 Flash retains the integer thinkingBudget ladder.""" + cases = { + "minimal": 512, + "low": 2048, + "medium": 8192, + "high": 24576, + "max": -1, + "xhigh": -1, + } + for effort, expected in cases.items(): + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + reasoning_effort = effort, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingBudget": expected}, (effort, tc) + + +def test_gemini3_flash_effort_levels_map_to_thinking_level(monkeypatch): + """Gemini 3 Flash thinkingLevel ladder: minimal/low/medium/high.""" + cases = { + "minimal": "minimal", + "low": "low", + "medium": "medium", + "high": "high", + "max": "high", + } + for effort, expected in cases.items(): + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + reasoning_effort = effort, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": expected}, (effort, tc) + + +def test_gemini3_pro_passes_medium_through(monkeypatch): + """Gemini 3.1+ Pro accepts thinkingLevel="medium" per + https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro; + forward as-is (medium is the documented mid-tier on Gemini 3.1).""" + for model in ( + "gemini-3.1-pro-preview", + "gemini-pro-latest", + ): + captured = _capture_body( + monkeypatch, + model = model, + reasoning_effort = "medium", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "medium"}, (model, tc) + + +def test_gemini3_pro_minimal_effort_coerces_to_low(monkeypatch): + """Gemini 3 Pro rejects thinkingLevel="minimal"; coerce to "low".""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.1-pro-preview", + reasoning_effort = "minimal", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, tc + + +def test_gemini3_flash_effort_none_maps_to_minimal(monkeypatch): + """reasoning_effort='none' on Gemini 3 Flash -> thinkingLevel=minimal.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + reasoning_effort = "none", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "minimal"}, tc + + +def test_thinking_default_omits_thinking_config(monkeypatch): + """When neither knob is supplied, thinkingConfig is omitted (Google's + server-side default applies).""" + captured = _capture_body(monkeypatch, model = "gemini-3.5-flash") + gc = captured["body"]["generationConfig"] + assert "thinkingConfig" not in gc, gc + + +def test_nano_banana_alias_routes_through_image_modalities(monkeypatch): + """`nano-banana-pro-preview` aliases the Pro image model; must set + responseModalities=[TEXT,IMAGE] when the Images pill is on + (enabled_tools includes "image_generation").""" + captured = _capture_body( + monkeypatch, + model = "nano-banana-pro-preview", + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert gc.get("responseModalities") == ["TEXT", "IMAGE"], gc + + +def test_image_capable_model_without_image_pill_stays_text_only(monkeypatch): + """When the Images pill is off (no image_generation in enabled_tools), an + image-capable model id (gemini-2.5-flash-image) must force + responseModalities=["TEXT"]. Google's image models default to text+image + when responseModalities is omitted, so omitting it would silently bill + image output the UI says is disabled.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = [], + ) + gc = captured["body"]["generationConfig"] + assert gc.get("responseModalities") == ["TEXT"], gc + + +def test_image_models_skip_thinking_config(monkeypatch): + """Image-tier ids have no visible thinking knob and must NOT forward + thinkingConfig even when stale UI state still sends `reasoning_effort` or + `enable_thinking=False`.""" + for model in ( + "gemini-2.5-flash-image", + "gemini-3.1-flash-image-preview", + "gemini-3-pro-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + reasoning_effort = "high", + enable_thinking = False, + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert "thinkingConfig" not in gc, (model, gc) + + +def test_image_models_drop_code_execution(monkeypatch): + """All image-tier ids reject `tools: [{codeExecution: {}}]`; drop + silently. (Gemini 3 image models DO accept googleSearch -- see + test_gemini3_image_models_allow_google_search; older ones drop + everything.)""" + for model in ( + "gemini-2.5-flash-image", + "gemini-3.1-flash-image-preview", + "gemini-3-pro-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + enabled_tools = ["image_generation", "code_execution"], + ) + tools_arr = captured["body"].get("tools") or [] + names = [list(t.keys())[0] for t in tools_arr] + assert "codeExecution" not in names, (model, tools_arr) + + +def test_gemini_35_pro_uses_thinking_level(monkeypatch): + """`gemini-3.5-pro` is Gemini 3 family and uses thinkingLevel (not + thinkingBudget). "Off" maps to "low" since Pro tier rejects "minimal".""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-pro", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, tc + + +def test_gemini3_image_models_allow_google_search(monkeypatch): + """Google documents Search grounding on the Gemini 3 image family + (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + nano-banana-pro). codeExecution stays blocked on image mode.""" + for model in ( + "gemini-3-pro-image-preview", + "gemini-3.1-flash-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + tools_arr = captured["body"].get("tools") or [] + names = [list(t.keys())[0] for t in tools_arr] + assert "googleSearch" in names, (model, tools_arr) + assert "codeExecution" not in names, (model, tools_arr) + + +def test_legacy_image_models_block_google_search(monkeypatch): + """Older Gemini image ids (gemini-2.5-flash-image) still 400 on + `tools: [{googleSearch: {}}]`; backend keeps stripping it.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + assert "tools" not in captured["body"], captured["body"].get("tools") + + +def test_legacy_openai_base_url_normalized(monkeypatch): + """Saved Gemini providers with the legacy `/v1beta/openai` base (from + pre-PR OpenAI-compat plumbing) now point at the native endpoint without + the user re-saving the connection.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://generativelanguage.googleapis.com/v1beta" + + +def test_finish_reason_swaps_to_tool_calls_when_function_call_emitted(monkeypatch): + """Gemini emits finishReason="STOP" even for pure functionCall turns; + surface as `tool_calls` so OAI clients run the tool.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "lookup", "args": {"k": "v"}}}], + }, + "finishReason": "STOP", + } + ] + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + finish_chunks = [ + c for c in chunks if c.get("choices", [{}])[0].get("finish_reason") is not None + ] + assert finish_chunks, chunks + assert finish_chunks[-1]["choices"][0]["finish_reason"] == "tool_calls", chunks + + +def test_thought_signature_round_trips_into_gemini_function_call(monkeypatch): + """An assistant tool_call carrying `extra_content.google.thought_signature` + must echo it back as a sibling of the Gemini functionCall part.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "lookup x"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "extra_content": {"google": {"thought_signature": "SIG-ABC"}}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_0", + "name": "lookup", + "content": "{}", + }, + ], + ) + contents = captured["body"]["contents"] + fc_turn = next((c for c in contents if c["role"] == "model"), None) + assert fc_turn is not None, contents + fc_part = next( + (p for p in fc_turn["parts"] if "functionCall" in p), + None, + ) + assert fc_part is not None, fc_turn + assert fc_part.get("thoughtSignature") == "SIG-ABC", fc_part + + +def test_thought_signature_emitted_in_tool_call_delta(monkeypatch): + """A Gemini functionCall part with `thoughtSignature` must surface it on + the outbound OpenAI tool_calls delta via + `extra_content.google.thought_signature`.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "lookup", + "args": {"k": "v"}, + "id": "call_xyz", + }, + "thoughtSignature": "SIG-FROM-GEMINI", + } + ], + }, + "finishReason": "STOP", + } + ] + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + deltas = [ + tc + for c in chunks + for tc in (c.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) + ] + assert deltas, chunks + sig = deltas[0].get("extra_content", {}).get("google", {}).get("thought_signature") + assert sig == "SIG-FROM-GEMINI", deltas + + +def test_image_models_suppress_phantom_web_search_card(monkeypatch): + """When the image guard filters googleSearch out of the request, the + inbound stream must NOT emit web_search tool_start / tool_end (else the UI + shows a misleading 'Search complete' card on a turn Gemini never + searched).""" + sse = [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "drawn"}]}, + "finishReason": "STOP", + } + ] + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + chunks = _parse_chunks(lines) + tool_evs = [ + ev + for c in chunks + for ev in [c.get("_toolEvent")] + if isinstance(ev, dict) and ev.get("tool_name") == "web_search" + ] + assert tool_evs == [], tool_evs + + +def test_image_generation_tool_on_image_model_drops_text_tools(monkeypatch): + """`enabled_tools=["image_generation", "web_search", "code_execution"]` + on a Gemini IMAGE model flips responseModalities to TEXT+IMAGE; in that + mode codeExecution must NOT be forwarded (Gemini rejects text code tools + alongside image responseModalities). Older image families also drop + googleSearch.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = [ + "image_generation", + "web_search", + "code_execution", + ], + ) + assert "tools" not in captured["body"], captured["body"] + assert captured["body"]["generationConfig"].get("responseModalities") == ["TEXT", "IMAGE"] + + +def test_prompt_feedback_block_reason_surfaces_as_error(monkeypatch): + """`promptFeedback.blockReason` with zero candidates must produce an error + chunk, not a silent empty assistant reply.""" + sse = [ + { + "promptFeedback": {"blockReason": "SAFETY"}, + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + error_chunks = [c for c in chunks if "error" in c] + assert error_chunks, chunks + assert "SAFETY" in (error_chunks[0].get("error", {}).get("message") or ""), error_chunks + + +def test_usage_chunk_includes_thoughts_tokens(monkeypatch): + """`thoughtsTokenCount` is the hidden-reasoning slice of output; roll it + into `output_tokens` AND surface it on + `output_tokens_details.reasoning_tokens` so total_tokens reflects the full + billable spend.""" + sse = [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 20, + "totalTokenCount": 35, + }, + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + usage_chunk = next((c for c in chunks if isinstance(c.get("usage"), dict)), None) + assert usage_chunk is not None, chunks + usage = usage_chunk["usage"] + assert usage.get("prompt_tokens") == 10, usage + # candidates 5 + thoughts 20 = 25 output tokens; total = 35. + assert usage.get("completion_tokens") == 25, usage + assert usage.get("total_tokens") == 35, usage + + +# ── web_search forwarded as googleSearch tool ──────────────────────── + + +def test_web_search_forwarded_as_google_search_tool(monkeypatch): + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search"], + ) + tools = captured["body"].get("tools") or [] + assert {"googleSearch": {}} in tools, tools + + +def test_code_execution_forwarded_as_code_execution_tool(monkeypatch): + captured = _capture_body( + monkeypatch, + enabled_tools = ["code_execution"], + ) + tools = captured["body"].get("tools") or [] + assert {"codeExecution": {}} in tools, tools + + +def test_omitted_tools_leaves_body_untouched(monkeypatch): + captured = _capture_body(monkeypatch, enabled_tools = []) + assert "tools" not in captured["body"], captured["body"] + + +# ── prompt caching passthrough ─────────────────────────────────────── + + +def test_cached_content_pass_through(monkeypatch): + """A string cache id on enable_prompt_caching is forwarded verbatim.""" + cache_name = "cachedContents/abc123" + captured = _capture_body( + monkeypatch, + enable_prompt_caching = cache_name, + ) + assert captured["body"].get("cachedContent") == cache_name + + +def test_boolean_caching_does_not_set_cached_content(monkeypatch): + """Studio's existing True/False signals shouldn't fabricate a cache id.""" + captured = _capture_body(monkeypatch, enable_prompt_caching = True) + assert "cachedContent" not in captured["body"] + + +# ── image generation: request modalities + response translation ────── + + +def test_image_model_sets_response_modalities(monkeypatch): + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + ) + assert captured["body"]["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] + + +def test_image_generation_tool_sets_response_modalities_on_image_model(monkeypatch): + """`enabled_tools=["image_generation"]` flips responseModalities + only when the selected model is image-capable; otherwise the + request stays plain text (text-only models 400 on + responseModalities).""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + ) + assert captured["body"]["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] + + +def test_image_response_emits_image_b64_tool_event(monkeypatch): + """`inlineData` parts become a tool_end with image_b64 + image_mime.""" + fake_b64 = base64.b64encode(b"PNG-BYTES").decode() + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": fake_b64, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 0, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + image_starts = [e for e in starts if e.get("tool_name") == "image_generation"] + image_ends = [e for e in ends if e.get("image_b64")] + assert len(image_starts) == 1, tool_events + assert len(image_ends) == 1, tool_events + assert image_ends[0]["image_b64"] == fake_b64 + assert image_ends[0]["image_mime"] == "image/png" + + +# ── function calling round-trips both directions ───────────────────── + + +def test_function_call_response_translates_to_tool_calls_delta(monkeypatch): + """Gemini `functionCall` parts become OpenAI `tool_calls` delta chunks.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "get_weather", + "args": {"location": "Paris"}, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 12, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + tool_call_chunks = [ + c + for c in chunks + if "_toolEvent" not in c + and any( + (isinstance(ch.get("delta"), dict) and "tool_calls" in ch["delta"]) + for ch in c.get("choices", []) + ) + ] + assert len(tool_call_chunks) == 1, chunks + tc = tool_call_chunks[0]["choices"][0]["delta"]["tool_calls"][0] + assert tc["function"]["name"] == "get_weather" + args = json.loads(tc["function"]["arguments"]) + assert args == {"location": "Paris"} + + +def test_tool_message_translates_to_function_response_part(monkeypatch): + """role=tool follow-ups are rewritten to functionResponse parts.""" + messages = [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "Paris"}), + }, + } + ], + }, + { + "role": "tool", + "name": "get_weather", + "content": json.dumps({"temp_c": 18, "summary": "Sunny"}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + # Last turn must be a functionResponse part (Gemini wraps it as a role=user + # turn carrying the result). + last = contents[-1] + assert last["role"] == "user", last + fr = last["parts"][0].get("functionResponse") + assert fr is not None, last + assert fr["name"] == "get_weather" + assert fr["response"] == {"temp_c": 18, "summary": "Sunny"} + # And the assistant turn carries the original functionCall so the model + # sees the round-trip context. + assistant_turn = [c for c in contents if c["role"] == "model"][0] + fc_part = next( + (p for p in assistant_turn["parts"] if "functionCall" in p), + None, + ) + assert fc_part is not None, assistant_turn + assert fc_part["functionCall"]["name"] == "get_weather" + assert fc_part["functionCall"]["args"] == {"location": "Paris"} + + +def test_parallel_function_calls_get_distinct_tool_call_indices(monkeypatch): + """Each emitted functionCall in one assistant turn needs its own + tool_calls[*].index. Hardcoding index=0 collapses parallel calls onto one + slot in OpenAI-style reassemblers.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "id": "call_alpha", + "name": "search", + "args": {"q": "alpha"}, + } + }, + { + "functionCall": { + "id": "call_beta", + "name": "search", + "args": {"q": "beta"}, + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + tool_call_chunks = [ + c + for c in chunks + if "_toolEvent" not in c + and any( + (isinstance(ch.get("delta"), dict) and "tool_calls" in ch["delta"]) + for ch in c.get("choices", []) + ) + ] + assert len(tool_call_chunks) == 2, tool_call_chunks + indices = [c["choices"][0]["delta"]["tool_calls"][0]["index"] for c in tool_call_chunks] + assert indices == [0, 1], indices + + +def test_function_call_ids_forwarded_into_gemini_function_call_part(monkeypatch): + """OpenAI tool_call id rides functionCall.id so parallel calls disambiguate.""" + messages = [ + {"role": "user", "content": "x"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_alpha", + "type": "function", + "function": { + "name": "search", + "arguments": json.dumps({"q": "a"}), + }, + }, + { + "id": "call_beta", + "type": "function", + "function": { + "name": "search", + "arguments": json.dumps({"q": "b"}), + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_alpha", + "content": json.dumps({"hits": ["A"]}), + }, + { + "role": "tool", + "tool_call_id": "call_beta", + "content": json.dumps({"hits": ["B"]}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + assistant_parts = next(c for c in contents if c["role"] == "model")["parts"] + call_ids = [p["functionCall"]["id"] for p in assistant_parts if "functionCall" in p] + assert call_ids == ["call_alpha", "call_beta"], assistant_parts + response_ids = [ + p["functionResponse"]["id"] for c in contents for p in c["parts"] if "functionResponse" in p + ] + assert response_ids == ["call_alpha", "call_beta"], contents + + +def test_parse_gemini_models_translates_native_catalog(): + """Gemini's native /v1beta/models payload becomes OpenAI-shape entries.""" + payload = { + "models": [ + { + "name": "models/gemini-2.5-flash", + "baseModelId": "gemini-2.5-flash", + "displayName": "Gemini 2.5 Flash", + "supportedGenerationMethods": [ + "generateContent", + "streamGenerateContent", + ], + }, + { + "name": "models/embedding-001", + "supportedGenerationMethods": ["embedContent"], + }, + { + "name": "models/gemini-2.5-pro", + }, + ] + } + out = ExternalProviderClient._parse_gemini_models(payload) + ids = [m["id"] for m in out] + assert "gemini-2.5-flash" in ids + assert "gemini-2.5-pro" in ids + assert "embedding-001" not in ids + flash = next(m for m in out if m["id"] == "gemini-2.5-flash") + assert flash["display_name"] == "Gemini 2.5 Flash" + assert flash["owned_by"] == "google" + + +def test_code_execution_parts_translate_to_code_execution_tool_events(monkeypatch): + """executableCode + codeExecutionResult parts emit code_execution events.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "print(2+2)", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "4\n", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse, enabled_tools = ["code_execution"]) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + code_starts = [ + e + for e in tool_events + if e.get("type") == "tool_start" and e.get("tool_name") == "code_execution" + ] + code_ends = [ + e for e in tool_events if e.get("type") == "tool_end" and "4" in str(e.get("result", "")) + ] + assert len(code_starts) == 1, tool_events + assert code_starts[0]["arguments"]["code"] == "print(2+2)" + assert code_starts[0]["arguments"]["language"] == "python" + assert len(code_ends) == 1, tool_events + # tool_start and tool_end must share a tool_call_id so the frontend pairs + # them onto one CodeExecutionToolUI block. + assert code_starts[0]["tool_call_id"] == code_ends[0]["tool_call_id"] + + +def test_code_execution_failure_outcome_surfaces_in_result(monkeypatch): + """OUTCOME_FAILED is prefixed onto the result text so the UI shows it.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "1/0", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_FAILED", + "output": "ZeroDivisionError", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 2, + }, + } + ] + lines = _collect(monkeypatch, sse, enabled_tools = ["code_execution"]) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + result_text = next( + (e["result"] for e in tool_events if e.get("type") == "tool_end"), + "", + ) + assert "OUTCOME_FAILED" in result_text + assert "ZeroDivisionError" in result_text + + +def test_tool_message_recovers_name_from_tool_call_id(monkeypatch): + """When name is omitted, recover it from the matching tool_call_id.""" + messages = [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_xyz", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "Paris"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_xyz", + "content": json.dumps({"temp_c": 18}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + last = contents[-1] + fr = last["parts"][0].get("functionResponse") + assert fr is not None, last + assert ( + fr["name"] == "get_weather" + ), "name should fall back to the prior tool_call's function name" + + +# ── usage chunk surfaces promptTokenCount / candidatesTokenCount ───── + + +def test_usage_chunk_translates_gemini_token_counts(monkeypatch): + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1234, + "candidatesTokenCount": 56, + "cachedContentTokenCount": 1000, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("choices") == [] and "usage" in c] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["prompt_tokens"] == 1234 + assert usage["completion_tokens"] == 56 + assert usage["total_tokens"] == 1290 + assert usage["prompt_tokens_details"]["cached_tokens"] == 1000 + + +# ── multimodal: vision image -> inlineData ─────────────────────────── + + +def test_vision_data_url_translates_to_inline_data(monkeypatch): + fake = base64.b64encode(b"JPGBYTES").decode() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{fake}", + }, + }, + ], + } + ] + captured = _capture_body(monkeypatch, messages = messages) + parts = captured["body"]["contents"][0]["parts"] + inline_parts = [p for p in parts if "inlineData" in p] + assert len(inline_parts) == 1, parts + assert inline_parts[0]["inlineData"] == {"mimeType": "image/jpeg", "data": fake} + + +# ── finish reason mapping ──────────────────────────────────────────── + + +@pytest.mark.parametrize( + "gemini_reason, openai_reason", + [ + ("STOP", "stop"), + ("MAX_TOKENS", "length"), + ("SAFETY", "content_filter"), + ("PROHIBITED_CONTENT", "content_filter"), + ], +) +def test_finish_reason_translation(monkeypatch, gemini_reason, openai_reason): + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "x"}], + }, + "finishReason": gemini_reason, + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + finish_chunks = [ + c for c in chunks if any(ch.get("finish_reason") for ch in c.get("choices", [])) + ] + assert any( + ch["choices"][0]["finish_reason"] == openai_reason for ch in finish_chunks + ), finish_chunks + + +# ── grounding citations surface as web_search tool_end ─────────────── + + +def test_grounding_metadata_surfaces_as_tool_end_citations(monkeypatch): + """`groundingMetadata.groundingChunks[].web` -> tool_end result block.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Answer with sources."}], + }, + "groundingMetadata": { + "groundingChunks": [ + { + "web": { + "uri": "https://example.com/a", + "title": "Example A", + } + }, + { + "web": { + "uri": "https://example.com/b", + "title": "Example B", + } + }, + ] + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 7, + "candidatesTokenCount": 3, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["web_search"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + web_search_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "gemini_web_search" + ] + assert len(web_search_ends) == 1, tool_events + result = web_search_ends[0]["result"] + assert "https://example.com/a" in result + assert "https://example.com/b" in result + assert "Example A" in result + assert "Example B" in result + + +# ── round 3 review follow-ups ───────────────────────────────────────── + + +def test_custom_gemini_proxy_base_url_not_rewritten(): + """Only the Google-hosted /v1beta/openai base is normalized; a custom + gateway whose path ends in /openai must be left alone.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://proxy.example.com/team/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://proxy.example.com/team/openai" + + +def test_custom_gemini_proxy_uses_openai_dispatch(): + """Any non-Google Gemini base (LiteLLM, custom OpenAI-compat routers) must + route through the OpenAI-compatible forwarder, not the native translator. + Auth uses Authorization: Bearer ..., not x-goog-api-key.""" + for base in ( + "https://proxy.example.com/team/openai", + "https://proxy.example.com/v1", + "https://litellm.internal.example/v1", + ): + client = ExternalProviderClient( + provider_type = "gemini", + base_url = base, + api_key = "AIza-test-key", + ) + assert client._is_openai_compatible() is True, base + headers = client._auth_headers() + assert "x-goog-api-key" not in {k.lower() for k in headers}, (base, headers) + assert headers["Authorization"] == "Bearer AIza-test-key", (base, headers) + + +def test_google_hosted_gemini_still_uses_native_dispatch(): + """Google-hosted Gemini keeps native dispatch + x-goog-api-key auth.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + api_key = "AIza-test-key", + ) + assert client._is_openai_compatible() is False + headers = client._auth_headers() + assert headers.get("x-goog-api-key") == "AIza-test-key", headers + + +def test_invalid_gemini_model_id_rejected_before_request(monkeypatch): + """Path-traversal model ids must be rejected before the URL is + interpolated, so the configured API key isn't sent to unintended Gemini + endpoints.""" + + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response( + 200, + content = _gemini_sse([]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + out: list[str] = [] + + async def run(): + client = _make_gemini_client() + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "../cachedContents/leak", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + ): + out.append(line) + await client.close() + + _drive(run()) + # No outbound request should have been issued. + assert captured == [], captured + error_lines = [line for line in out if '"error"' in line] + assert error_lines, out + + +def test_top_k_omitted_when_not_explicit_default_for_gemini(monkeypatch): + """top_k=None means "use provider default"; helper must not emit `topK` in + generationConfig when the caller didn't pass it.""" + captured = _capture_body(monkeypatch, top_k = None) + assert "topK" not in captured["body"]["generationConfig"], captured["body"] + + +def test_text_model_image_generation_tool_silently_dropped(monkeypatch): + """A stale `enabled_tools=["image_generation"]` on a text-only Gemini + model (e.g. gemini-2.5-flash) must NOT switch the request into image mode + -- Google's API 400s on responseModalities for text models.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert "responseModalities" not in gc, gc + + +def test_empty_text_part_with_thought_signature_emits_extra_content(monkeypatch): + """Gemini 3 can ship a content-free fragment whose only payload is + `thoughtSignature`. The translator must still surface it on a + delta.extra_content envelope so the next turn can replay it.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"text": "answer"}, + {"thoughtSignature": "SIG-FINAL"}, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 2, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + extra_carriers = [ + c + for c in chunks + if c.get("choices") + and c["choices"][0]["delta"].get("extra_content") + == {"google": {"thought_signature": "SIG-FINAL"}} + ] + assert extra_carriers, chunks + + +def test_enable_prompt_caching_false_string_coerces_to_bool(): + """Pre-PR the field was Optional[bool]; widening to Union[bool,str] must + preserve historical coercion so callers sending `"false"` still opt out of + caching.""" + from models.inference import ChatCompletionRequest + + msg = {"role": "user", "content": "hi"} + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "false", + } + ) + assert req.enable_prompt_caching is False, req.enable_prompt_caching + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "true", + } + ) + assert req.enable_prompt_caching is True + + # An actual cache resource name passes through untouched. + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "cachedContents/abc123", + } + ) + assert req.enable_prompt_caching == "cachedContents/abc123" + + +def test_legacy_google_openai_base_url_is_rewritten(): + """The Google-hosted /v1beta/openai legacy base IS still rewritten.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://generativelanguage.googleapis.com/v1beta" + + +def test_remote_image_url_downloads_and_inlines_as_base64(monkeypatch): + """Round 14: arbitrary public HTTPS image URLs cannot be sent as Gemini + fileData (reserved for Files API URIs and YouTube). The translator must + fetch the bytes server-side and inline them as base64 inlineData.""" + image_bytes = b"FAKEPNGBYTES" + + async def fake_fetch( + url, + fallback_mime, + max_bytes = None, + ): + assert url == "https://cdn.example.com/diagram.png" + return ("image/png", base64.b64encode(image_bytes).decode("ascii")) + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": { + "url": "https://cdn.example.com/diagram.png", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + inline = next((p for p in parts if "inlineData" in p), None) + assert inline is not None, parts + assert inline["inlineData"]["mimeType"] == "image/png" + assert inline["inlineData"]["data"] == base64.b64encode(image_bytes).decode() + assert not any("fileData" in p for p in parts), parts + + +def test_remote_image_url_dropped_when_fetch_returns_none(monkeypatch): + """Round 15: if the SSRF guard rejects the URL (private host, non-https, + oversize, non-image), the helper returns None and the image part is + silently dropped, not forwarded as raw bytes or a fileData fallback.""" + + async def fake_fetch_reject( + url, + fallback_mime, + max_bytes = None, + ): + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch_reject) + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": {"url": "http://10.0.0.5/private.png"}, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + assert not any("inlineData" in p for p in parts), parts + assert not any("fileData" in p for p in parts), parts + + +def test_safe_fetch_image_rejects_non_https(): + """SSRF guard: only https URLs may be fetched.""" + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini("http://cdn.example.com/x.png", "image/png") + ) + assert res is None + + +def test_safe_fetch_image_rejects_loopback_ip_literal(): + """SSRF guard: refuse loopback / private IP literals before any network + call.""" + for url in ( + "https://127.0.0.1/x.png", + "https://[::1]/x.png", + "https://169.254.169.254/latest/meta-data", + "https://10.0.0.5/x.png", + "https://192.168.1.1/x.png", + ): + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini(url, "image/png") + ) + assert res is None, url + + +def test_safe_fetch_image_rejects_resolved_private_host(monkeypatch): + """SSRF guard: if a hostname resolves to a private IP, refuse.""" + import socket + + def fake_getaddrinfo(host, *_args, **_kwargs): + return [(socket.AF_INET, None, None, "", ("10.0.0.5", 0))] + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini("https://internal.example/x.png", "image/png") + ) + assert res is None + + +def test_youtube_and_files_api_uris_stay_as_file_data(monkeypatch): + """Round 14: YouTube URLs and generativelanguage.googleapis.com Files API + URIs are the documented `fileData.fileUri` paths and must NOT be + downloaded; arbitrary public URLs do get fetched.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "explain"}, + { + "type": "image_url", + "image_url": { + "url": "https://www.youtube.com/watch?v=abc123", + }, + }, + { + "type": "image_url", + "image_url": { + "url": "https://generativelanguage.googleapis.com/v1beta/files/abc", + }, + }, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + parts = captured["body"]["contents"][-1]["parts"] + file_uris = [p["fileData"]["fileUri"] for p in parts if "fileData" in p] + assert "https://www.youtube.com/watch?v=abc123" in file_uris, parts + assert "https://generativelanguage.googleapis.com/v1beta/files/abc" in file_uris, parts + + +def test_tool_use_prompt_tokens_added_to_input_tokens(monkeypatch): + """`toolUsePromptTokenCount` must roll into the OpenAI prompt total -- + else tool turns silently undercount input tokens.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "result"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "toolUsePromptTokenCount": 100, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 2, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("usage")] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["prompt_tokens"] == 110, usage + assert usage["completion_tokens"] == 7, usage + assert usage["total_tokens"] == 117, usage + assert usage["completion_tokens_details"]["reasoning_tokens"] == 2, usage + + +def test_usage_chunk_reasoning_tokens_surfaced(monkeypatch): + """thoughtsTokenCount must surface as + completion_tokens_details.reasoning_tokens in the emitted OpenAI usage + chunk.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 20, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("usage")] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["completion_tokens"] == 25, usage + assert usage["completion_tokens_details"]["reasoning_tokens"] == 20, usage + + +def test_prompt_block_pairs_web_search_tool_end(monkeypatch): + """When `promptFeedback.blockReason` triggers after the synthetic + web_search tool_start, the helper must emit a matching tool_end so the UI + doesn't leave a "searching..." spinner stuck on screen.""" + sse = [ + {"promptFeedback": {"blockReason": "SAFETY"}}, + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["web_search"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + assert len(starts) == 1, tool_events + assert len(ends) == 1, tool_events + assert ends[0]["tool_call_id"] == "gemini_web_search" + assert "aborted" in ends[0]["result"] + error_chunks = [c for c in chunks if c.get("error")] + assert error_chunks, chunks + + +def test_code_execution_tool_events_stow_native_part(monkeypatch): + """executableCode / codeExecutionResult must round-trip native ids and + thoughtSignature in google.native_part so follow-up turns can replay + Gemini's required history shape.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(1+1)", + }, + "thoughtSignature": "SIG-CODE", + }, + { + "codeExecutionResult": { + "id": "result_a", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + code_start = next( + (e for e in starts if e.get("tool_name") == "code_execution"), + None, + ) + code_end = next(iter(ends), None) + assert code_start is not None, starts + assert code_start["tool_call_id"] == "code_a", code_start + native = code_start["arguments"]["google"]["native_part"] + # Round 21: native_part uses an ordered `parts` list so per-part + # `thoughtSignature` survives a frontend merge of executableCode + + # codeExecutionResult into one tool-call card. + start_parts = native["parts"] + assert start_parts[0]["executableCode"]["id"] == "code_a" + assert start_parts[0]["thoughtSignature"] == "SIG-CODE" + assert code_end is not None, ends + assert code_end["tool_call_id"] == "code_a", code_end + native_end = code_end["google"]["native_part"] + end_parts = native_end["parts"] + assert end_parts[0]["codeExecutionResult"]["id"] == "result_a" + + +def test_inline_image_tool_end_carries_thought_signature(monkeypatch): + """Inline image parts with thoughtSignature must persist it on the emitted + tool_end so Gemini 3 image editing can echo it back.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": base64.b64encode(b"PNG").decode(), + }, + "thoughtSignature": "SIG-IMG", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 4, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + image_ends = [e for e in tool_events if e.get("type") == "tool_end" and e.get("image_b64")] + assert image_ends, tool_events + assert image_ends[0]["google"]["thought_signature"] == "SIG-IMG" + # Multi-turn image edit must replay the original inlineData part with its + # thoughtSignature; the outbound translator reads + # google.native_part.parts[].inlineData, so stow it on the tool_end too. + # Round 21 made native_part an ordered parts list so a per-part signature + # stays attached to inlineData only. + native = image_ends[0]["google"]["native_part"] + image_parts = native["parts"] + assert image_parts[0]["inlineData"]["mimeType"] == "image/png" + assert image_parts[0]["inlineData"]["data"] == base64.b64encode(b"PNG").decode() + assert image_parts[0]["thoughtSignature"] == "SIG-IMG" + + +def test_code_execution_plot_attaches_inline_image_native_part(monkeypatch): + """A code_execution turn that returns a matplotlib plot must stow the + plot's inlineData on the secondary tool_end so the follow-up turn can + replay the image alongside executableCode and codeExecutionResult.""" + plot_data = base64.b64encode(b"PLOT").decode() + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "plt.plot([0,1])", + }, + }, + { + "codeExecutionResult": { + "id": "result_a", + "outcome": "OUTCOME_OK", + "output": "", + }, + }, + { + "inlineData": { + "mimeType": "image/png", + "data": plot_data, + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + code_ends = [ + e for e in tool_events if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_a" + ] + # Two tool_end events on the same id: one for codeExecutionResult, one + # merging in the inlineData plot. The plot one must carry the native + # inlineData under google.native_part so the frontend tool_end merge union + # joins it with the prior executableCode and codeExecutionResult parts on + # the same card. + assert len(code_ends) == 2, code_ends + image_end = next( + (e for e in code_ends if "__IMAGES__:" in (e.get("result") or "")), + None, + ) + assert image_end is not None, code_ends + native = image_end["google"]["native_part"] + plot_parts = native["parts"] + assert plot_parts[0]["inlineData"]["mimeType"] == "image/png" + assert plot_parts[0]["inlineData"]["data"] == plot_data + + +def test_text_chunk_carries_thought_signature(monkeypatch): + """Text parts with thoughtSignature surface it on delta.extra_content so + frontend persistence can replay it on the follow-up turn.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "hello", + "thoughtSignature": "SIG-TEXT", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 2, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + text_chunks = [ + c for c in chunks if c.get("choices") and c["choices"][0]["delta"].get("content") == "hello" + ] + assert text_chunks, chunks + extra = text_chunks[0]["choices"][0]["delta"].get("extra_content") + assert extra == {"google": {"thought_signature": "SIG-TEXT"}}, text_chunks + + +def test_openai_tools_translated_into_function_declarations(monkeypatch): + """Standard ChatCompletionRequest.tools must be forwarded into Gemini's + tools[].functionDeclarations envelope.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + }, + "required": ["city"], + }, + }, + } + ], + tool_choice = {"type": "function", "function": {"name": "get_weather"}}, + ) + tools_arr = captured["body"].get("tools") or [] + fn_decls = [t for t in tools_arr if "functionDeclarations" in t] + assert fn_decls, captured["body"] + decls = fn_decls[0]["functionDeclarations"] + assert decls[0]["name"] == "get_weather" + assert decls[0]["parameters"]["properties"]["city"]["type"] == "string" + tool_config = captured["body"].get("toolConfig") + assert tool_config is not None, captured["body"] + fcc = tool_config["functionCallingConfig"] + assert fcc["mode"] == "ANY" + assert fcc["allowedFunctionNames"] == ["get_weather"] + + +def test_tool_choice_auto_maps_to_function_calling_mode_auto(monkeypatch): + """tool_choice="auto" maps to toolConfig.functionCallingConfig.mode.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": {"name": "noop", "parameters": {"type": "object"}}, + } + ], + tool_choice = "auto", + ) + fcc = captured["body"]["toolConfig"]["functionCallingConfig"] + assert fcc["mode"] == "AUTO" + assert "allowedFunctionNames" not in fcc + + +def test_code_exec_inline_image_attaches_to_code_execution_card(monkeypatch): + """A codeExecution sandbox plot (matplotlib) ships as an inline image part + right after the codeExecutionResult. Instead of a separate empty + image_generation card, attach to the same code_execution tool_end via the + `__IMAGES__:` marker the chat adapter already understands.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_plot", + "language": "PYTHON", + "code": "import matplotlib.pyplot as plt; plt.plot([1,2,3]); plt.savefig('out.png')", + }, + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "saved", + }, + }, + { + "inlineData": { + "mimeType": "image/png", + "data": base64.b64encode(b"PNGDATA").decode(), + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + # No standalone image_generation card should have been emitted. + image_starts = [ + e + for e in tool_events + if e.get("type") == "tool_start" and e.get("tool_name") == "image_generation" + ] + assert not image_starts, tool_events + # The code_execution tool_end should now carry the inline image + # via the `__IMAGES__:` marker. + code_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_plot" + ] + assert code_ends, tool_events + final_result = code_ends[-1]["result"] + assert "__IMAGES__:" in final_result, code_ends + assert "data:image/png;base64," in final_result, code_ends + + +def test_code_execution_tool_call_replays_native_executable_code(monkeypatch): + """An assistant tool_call with toolName=code_execution and + extra_content.google.native_part holding the originally-emitted + `executableCode` + `codeExecutionResult` must round-trip as native Gemini + parts (not a generic functionCall) on the next turn.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "compute 2+2"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "code_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(2+2)", + }, + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "4\n", + }, + "thoughtSignature": "SIG-CODE", + }, + }, + }, + }, + ], + }, + {"role": "user", "content": "what was that result"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + native_keys = [list(p.keys())[0] for p in parts if isinstance(p, dict)] + assert "executableCode" in native_keys, parts + assert "codeExecutionResult" in native_keys, parts + assert not any( + "functionCall" in p and (p["functionCall"] or {}).get("name") == "code_execution" + for p in parts + ), parts + exec_part = next(p for p in parts if "executableCode" in p) + assert exec_part.get("thoughtSignature") == "SIG-CODE", exec_part + + +def test_image_generation_tool_call_replays_native_inline_data(monkeypatch): + """An assistant tool_call with toolName=image_generation and + extra_content.google.native_part.inlineData must replay the prior image as + a native Gemini inlineData part (not a generic functionCall) so multi-turn + image editing keeps the image context.""" + pixel = base64.b64encode(b"PNG").decode() + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + messages = [ + {"role": "user", "content": "make a circle"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "img_a", + "type": "function", + "function": { + "name": "image_generation", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "inlineData": { + "mimeType": "image/png", + "data": pixel, + }, + "thoughtSignature": "SIG-IMG", + }, + }, + }, + }, + ], + }, + {"role": "user", "content": "now make it blue"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + inline_parts = [p for p in parts if "inlineData" in p] + assert inline_parts, parts + assert inline_parts[0]["inlineData"]["mimeType"] == "image/png" + assert inline_parts[0]["inlineData"]["data"] == pixel + assert inline_parts[0].get("thoughtSignature") == "SIG-IMG", inline_parts + assert not any( + "functionCall" in p and (p["functionCall"] or {}).get("name") == "image_generation" + for p in parts + ), parts + + +def test_assistant_text_thought_signature_replays_on_outbound_text_part(monkeypatch): + """Assistant text with extra_content.google.thought_signature must attach + `thoughtSignature` to the LAST text part of the replayed Gemini history. + Gemini 3 strict function-calling rejects history that drops returned + signatures, so the frontend stows the latest signed-text signature and the + backend pins it on the next turn.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "hello"}, + ], + "extra_content": { + "google": {"thought_signature": "SIG-TEXT"}, + }, + }, + {"role": "user", "content": "again"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + text_parts = [p for p in parts if "text" in p] + assert text_parts, parts + assert text_parts[-1].get("thoughtSignature") == "SIG-TEXT", text_parts + + +def test_function_declarations_strip_openai_only_schema_keys(monkeypatch): + """OpenAI strict tools commonly include `additionalProperties`, `$schema`, + `$defs`, `strict`, etc. Gemini's Schema rejects those with + INVALID_ARGUMENT, so the translator must strip them while keeping + properties..type intact.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a value.", + "parameters": { + "type": "object", + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": False, + "strict": True, + "properties": { + "key": { + "type": "string", + "additionalProperties": False, + }, + }, + "required": ["key"], + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + (t.get("functionDeclarations") for t in tools_arr if "functionDeclarations" in t), + None, + ) + assert decls is not None, captured["body"] + params = decls[0]["parameters"] + assert "additionalProperties" not in params + assert "$schema" not in params + assert "strict" not in params + assert params["type"] == "object" + assert params["properties"]["key"]["type"] == "string" + assert "additionalProperties" not in params["properties"]["key"] + assert params["required"] == ["key"] + + +def test_function_declarations_inline_local_refs_into_gemini_schema(monkeypatch): + """Round 25: Pydantic-generated tool schemas hoist nested object shapes + into `$defs` and reference them with `{"$ref": "#/$defs/..."}`. Gemini's + OpenAPI subset has no $ref, so a naive allowlist sanitizer drops the + reference and reduces the nested property to `{}`, losing its type, fields, + and required keys. The sanitizer must resolve local `#/...` pointers and + inline the referenced schema.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "set_user", + "description": "Persist a user.", + "parameters": { + "type": "object", + "$defs": { + "Address": { + "type": "object", + "properties": { + "street": {"type": "string"}, + "zip": {"type": "string"}, + }, + "required": ["street", "zip"], + }, + }, + "properties": { + "name": {"type": "string"}, + "address": {"$ref": "#/$defs/Address"}, + }, + "required": ["name", "address"], + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + (t.get("functionDeclarations") for t in tools_arr if "functionDeclarations" in t), + None, + ) + assert decls is not None, captured["body"] + params = decls[0]["parameters"] + assert "$defs" not in params + address = params["properties"]["address"] + assert address.get("type") == "object", address + assert address.get("properties", {}).get("street", {}).get("type") == "string" + assert address.get("properties", {}).get("zip", {}).get("type") == "string" + assert address.get("required") == ["street", "zip"] + + +def test_function_declarations_inline_local_refs_in_anyof_and_items(monkeypatch): + """The recursive inliner must reach through `anyOf` branches and `items` + (array element schemas), not just top-level property refs.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "bulk_set", + "parameters": { + "type": "object", + "$defs": { + "Address": { + "type": "object", + "properties": {"zip": {"type": "string"}}, + "required": ["zip"], + }, + }, + "properties": { + "primary": { + "anyOf": [ + {"$ref": "#/$defs/Address"}, + {"type": "null"}, + ], + }, + "extras": { + "type": "array", + "items": {"$ref": "#/$defs/Address"}, + }, + }, + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + (t.get("functionDeclarations") for t in tools_arr if "functionDeclarations" in t), + None, + ) + assert decls is not None + params = decls[0]["parameters"] + primary = params["properties"]["primary"] + # anyOf with single non-null branch + null collapses to inline + + # nullable: true; the inlined branch must contain the resolved Address + # shape. + assert primary.get("nullable") is True + assert primary.get("type") == "object" + assert primary.get("properties", {}).get("zip", {}).get("type") == "string" + extras = params["properties"]["extras"] + assert extras.get("type") == "array" + assert extras.get("items", {}).get("type") == "object" + assert extras.get("items", {}).get("properties", {}).get("zip", {}).get("type") == "string" + + +def test_function_declarations_self_referential_schema_terminates(monkeypatch): + """Self-referential / cyclic JSON Schemas (a `Node` with `children: + [Node]`) must not infinite-loop. The inliner tracks the set of refs in + flight and short-circuits to `{}` on a cycle.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "set_tree", + "parameters": { + "type": "object", + "$defs": { + "Node": { + "type": "object", + "properties": { + "value": {"type": "string"}, + "children": { + "type": "array", + "items": {"$ref": "#/$defs/Node"}, + }, + }, + }, + }, + "properties": { + "root": {"$ref": "#/$defs/Node"}, + }, + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + (t.get("functionDeclarations") for t in tools_arr if "functionDeclarations" in t), + None, + ) + assert decls is not None + root = decls[0]["parameters"]["properties"]["root"] + assert root.get("type") == "object" + assert root.get("properties", {}).get("value", {}).get("type") == "string" + + +def test_gemini_native_skips_orphan_function_response_for_dropped_builtin(monkeypatch): + """Round 26: when the assistant-side synthetic web_search/web_fetch + tool_call is dropped from native Gemini history, the matching role="tool" + follow-up must also be dropped. Otherwise the outbound body carries an + orphan functionResponse with no preceding functionCall, which 400s the + Gemini turn.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_s", + "type": "function", + "function": { + "name": "web_search", + "arguments": ('{"_server_tool": true, "query": "x"}'), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_s", + "content": "[search result]", + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + captured = _capture_body(monkeypatch, messages = built) + contents = captured["body"].get("contents") or [] + for entry in contents: + for part in entry.get("parts", []): + fr = part.get("functionResponse") + if isinstance(fr, dict): + assert fr.get("name") != "web_search", contents + + +def test_gemini_native_skips_orphan_function_response_for_native_part_replay(monkeypatch): + """Round 26: code_execution / image_generation tool_calls are replayed as + Gemini-native executableCode / codeExecutionResult / inlineData parts. The + matching role="tool" follow-up must NOT then be emitted as a + functionResponse named code_execution -- there is no declared user + function with that name, and Gemini's history rules already attribute the + result to the native parts above.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "plot something"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "print(2)", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "2\n", + } + }, + ] + } + } + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ], + ) + contents = captured["body"].get("contents") or [] + saw_native = False + for entry in contents: + for part in entry.get("parts", []): + if "executableCode" in part or "codeExecutionResult" in part: + saw_native = True + fr = part.get("functionResponse") + if isinstance(fr, dict): + assert fr.get("name") != "code_execution", contents + assert saw_native, contents + + +def test_gemini_native_part_falls_back_to_args_google(monkeypatch): + """Round 27: a direct OpenAI-compat API caller (or imported third-party + thread) cannot use Studio's non-standard `tool_calls[].extra_content` + field, so the native_part payload round-trips through `function.arguments` + as `{"google": {"native_part": {...}}}`. The synthetic-builtin detector + recognizes that location, but the replay branch was only reading from + `tc.extra_content.google.native_part`. Result: the round-25 guard saw a + synthetic builtin with no _native_part and dropped the entire assistant + turn, losing the prior code/image context. The translator must fall back + to args.google.native_part and still emit the native executableCode / + inlineData parts.""" + import json as _json + + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "draw a cat"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_img", + "type": "function", + "function": { + "name": "image_generation", + "arguments": _json.dumps( + { + "google": { + "native_part": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "AAAA", + } + } + ] + } + } + } + ), + }, + } + ], + }, + {"role": "user", "content": "now make it a dog"}, + ], + ) + contents = captured["body"].get("contents") or [] + saw_inline = False + for entry in contents: + for part in entry.get("parts", []): + if "inlineData" in part: + saw_inline = True + assert saw_inline, contents + + +def test_gemini_native_skips_synthetic_server_builtin_replay(monkeypatch): + """Round 25: Marked server-side builtin tool_calls (web_search / + web_fetch with `_server_tool` or `args.google.native_part`) must not fall + through to the generic Gemini `functionCall` replay path when no replayable + native part exists. Without this guard the outbound body contains a fake + `functionCall` whose name isn't a declared user function, and the Gemini + turn 400s.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_s", + "type": "function", + "function": { + "name": "web_search", + "arguments": ('{"_server_tool": true, "query": "x"}'), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_s", + "content": "[search result]", + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + captured = _capture_body(monkeypatch, messages = built) + contents = captured["body"].get("contents") or [] + for entry in contents: + for part in entry.get("parts", []): + fc = part.get("functionCall") + if isinstance(fc, dict): + assert fc.get("name") != "web_search", contents + + +def test_chat_message_extra_content_round_trips_through_validation(): + """Round 9: ChatMessage was missing `extra_content`, so Pydantic discarded + it during request validation and the text-part signature replay path read + nothing. The field must survive model_validate and pass through + _build_external_messages.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "hello"}, + ], + "extra_content": { + "google": {"thought_signature": "SIG-TEXT"}, + }, + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + assistant_msg = req.messages[1] + assert assistant_msg.extra_content == {"google": {"thought_signature": "SIG-TEXT"}} + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + assistant_out = built[1] + assert assistant_out["extra_content"] == {"google": {"thought_signature": "SIG-TEXT"}} + # Non-Gemini providers must NOT receive extra_content; Google's + # thought_signature is unknown to OpenAI / Mistral / etc. + built_openai = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + ) + assert "extra_content" not in built_openai[1], built_openai[1] + # Custom non-Google Gemini bases (LiteLLM / OAI-compat gateways) also must + # not receive Gemini-only extra_content -- the backend dispatches them + # through /chat/completions. + built_custom = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://litellm.example/v1", + ) + assert "extra_content" not in built_custom[1], built_custom[1] + + +def test_parallel_tool_results_group_into_one_user_block(monkeypatch): + """Round 14: Gemini docs group parallel functionResponses in a single + subsequent user content with multiple functionResponse parts. Consecutive + OpenAI role="tool" messages must merge into one Gemini user block, not + split into separate user turns.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "compute"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "add", "arguments": '{"x":1}'}, + }, + { + "id": "call_b", + "type": "function", + "function": {"name": "mul", "arguments": '{"x":2}'}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "add", + "content": "2", + }, + { + "role": "tool", + "tool_call_id": "call_b", + "name": "mul", + "content": "4", + }, + ], + ) + contents = captured["body"]["contents"] + # Initial user, model with two functionCalls, ONE user with two + # functionResponses. + tool_result_users = [ + c + for c in contents + if c.get("role") == "user" + and all(isinstance(p, dict) and "functionResponse" in p for p in (c.get("parts") or [])) + ] + assert len(tool_result_users) == 1, contents + fr_parts = tool_result_users[0]["parts"] + assert len(fr_parts) == 2, fr_parts + names = [p["functionResponse"]["name"] for p in fr_parts] + assert names == ["add", "mul"], names + + +def test_function_schema_nullable_type_array_flattens(monkeypatch): + """Round 14: OpenAI strict tools commonly use `"type": ["string", "null"]` + for optional fields. Gemini's OpenAPI-style Schema rejects union types and + expects `"type": "string"` with `"nullable": true`. The sanitizer must + translate the union form.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "city": {"type": ["string", "null"]}, + "score": {"type": ["number", "null"]}, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + params = decls[0]["parameters"]["properties"] + assert params["city"]["type"] == "string" + assert params["city"]["nullable"] is True + assert params["score"]["type"] == "number" + assert params["score"]["nullable"] is True + + +def test_image_picker_model_with_search_off_pill_strips_text_tools(monkeypatch): + """Round 11: image-tier model ids reject text-only tools and + thinkingConfig at the model level regardless of the Images pill. Selecting + gemini-2.5-flash-image + enabled_tools=["web_search"] with no + image_generation must NOT forward googleSearch or thinkingConfig (Gemini + 400s on text tools for legacy image ids).""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["web_search"], + reasoning_effort = "high", + ) + body = captured["body"] + assert "tools" not in body, body.get("tools") + assert "thinkingConfig" not in body.get("generationConfig", {}), body["generationConfig"] + + +def test_image_models_drop_function_declarations(monkeypatch): + """Image-mode requests cannot mix tools with responseModalities, so + user-supplied function declarations must be dropped.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tools = [ + { + "type": "function", + "function": {"name": "noop", "parameters": {"type": "object"}}, + } + ], + ) + assert captured["body"].get("tools") is None + assert captured["body"]["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] + + +def test_safe_fetch_image_rejects_malformed_bracketed_url(): + """Round 17: bracketed IPv6 garbage like `https://[bad/x.png` makes + urlparse raise ValueError. The fetch helper must catch it and drop the + image rather than crashing the request mid-build.""" + res = _drive(ep_mod._safe_fetch_image_for_gemini("https://[bad/x.png", "image/png")) + assert res is None + + +def test_safe_fetch_image_pins_validated_ip_no_hostname_in_request(monkeypatch): + """Round 17: the fetch helper must pin the validated IP into the outgoing + request URL (with a Host header carrying the original hostname). A second + hostname-style getaddrinfo after validate would be a DNS-rebinding gap, so + we assert the urllib opener is called with an IP-rewritten URL.""" + import socket + + captured: dict = {"requests": []} + + # Public IP during validate; record every getaddrinfo call. + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + captured.setdefault("dns", []).append(host) + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("8.8.8.8", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + headers = {"content-type": "image/png", "content-length": "3"} + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + return b"PNG" + + class _StubOpener: + def open( + self, + req, + timeout = None, + ): + captured["requests"].append( + { + "url": req.full_url, + "host_header": req.get_header("Host"), + } + ) + return _StubResp() + + monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + + res = _drive(ep_mod._safe_fetch_image_for_gemini("https://cdn.example.com/x.png", "image/png")) + assert res is not None + assert res[0] == "image/png" + # Outgoing URL must use the pinned IP literal, not the hostname. + assert any("8.8.8.8" in r["url"] for r in captured["requests"]), captured + assert all("cdn.example.com" not in r["url"] for r in captured["requests"]), captured + # Host header still carries the original hostname for vhost/SNI. + assert captured["requests"][0]["host_header"] == "cdn.example.com" + + +def test_safe_fetch_image_redirect_to_private_host_rejected(monkeypatch): + """Round 17: each redirect hop must re-validate the new host. A public hop + that redirects to an internal address must be dropped.""" + import socket + import urllib.error + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + if host == "internal.bad": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("10.0.0.5", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubOpener: + def open( + self, + req, + timeout = None, + ): + # Simulate a 302 to a private host. + raise urllib.error.HTTPError( + req.full_url, + 302, + "Found", + {"Location": "https://internal.bad/secret.png"}, + None, + ) + + monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + + res = _drive(ep_mod._safe_fetch_image_for_gemini("https://cdn.example.com/x.png", "image/png")) + assert res is None + + +def test_files_api_substring_url_not_misclassified_as_filedata(monkeypatch): + """Round 17: a CDN URL whose path/query merely contains the Files API + substring must NOT be sent as `fileData.fileUri`; route it through the + safe-fetch path. The old substring check + `"generativelanguage.googleapis.com/" in url.lower()` matched any URL + carrying that text anywhere.""" + captured_outbound: dict = {} + fetch_calls: list[str] = [] + + async def fake_fetch( + url, + fallback_mime, + max_bytes = None, + ): + fetch_calls.append(url) + return "image/png", base64.b64encode(b"DATA").decode("ascii") + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + captured_outbound["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + { + "type": "image_url", + "image_url": { + # Files-API-looking path, but host is an + # attacker CDN. + "url": "https://evil.example/path/generativelanguage.googleapis.com/v1beta/files/abc.png", + }, + }, + { + "type": "image_url", + "image_url": { + # Looks YouTube-ish in the path. + "url": "https://cdn.example.com/youtube.com/cat.png", + }, + }, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + parts = captured_outbound["body"]["contents"][-1]["parts"] + assert not any("fileData" in p for p in parts), parts + inline_count = sum(1 for p in parts if "inlineData" in p) + assert inline_count == 2, parts + assert len(fetch_calls) == 2, fetch_calls + + +def test_function_schema_anyof_null_variant_flattens_to_nullable(monkeypatch): + """Round 17: OpenAI/Pydantic emit `anyOf: [{X}, {"type":"null"}]` for + Optional[X]. Gemini's OpenAPI subset rejects `"type":"null"` inside anyOf. + The sanitizer must collapse a singleton-plus-null union back to the + non-null branch with `nullable: true`.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "label": { + "anyOf": [ + {"type": "string"}, + {"type": "null"}, + ] + }, + "count": { + "anyOf": [ + {"type": "integer"}, + {"type": "null"}, + ] + }, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + params = decls[0]["parameters"]["properties"] + assert params["label"]["type"] == "string" + assert params["label"]["nullable"] is True + assert "anyOf" not in params["label"] + assert params["count"]["type"] == "integer" + assert params["count"]["nullable"] is True + + +def test_legacy_gemini3_pro_medium_coerced_to_high(monkeypatch): + """Round 17: legacy `gemini-3-pro*` (incl. `-preview`, shut down + 2026-03-09) only accepted low/high. 3.1+ Pro added medium. The backend + must coerce medium → high for the legacy model so stale UI state doesn't + 400 the request.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3-pro-preview", + reasoning_effort = "medium", + ) + assert captured["body"]["generationConfig"]["thinkingConfig"] == {"thinkingLevel": "high"} + + +def test_gemini_3_1_pro_medium_passes_through(monkeypatch): + """Round 17 regression: 3.1+ Pro accepts medium; coercion must NOT apply + when the model id is gemini-3.1-pro*.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.1-pro-preview", + reasoning_effort = "medium", + ) + assert captured["body"]["generationConfig"]["thinkingConfig"] == {"thinkingLevel": "medium"} + + +def test_tool_calls_extra_content_stripped_for_non_native_gemini(): + """Round 17: per-tool-call `extra_content` (Gemini thoughtSignature + carrier) must not leak through `_build_external_messages` to + non-native-Gemini providers; OpenAI / Anthropic / custom Gemini OAI-compat + gateways would 400 on the unknown key.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "extra_content": { + "google": {"thought_signature": "SIG"}, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + + # Non-native providers (openai, custom Gemini OAI-compat proxy) must have + # extra_content stripped from the tool_call entry. + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + assert len(result) == 1 + tc = result[0]["tool_calls"][0] + assert "extra_content" not in tc, (provider_type, tc) + + # Native Gemini still receives extra_content for the round-trip. + result_native = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + tc_native = result_native[0]["tool_calls"][0] + assert tc_native["extra_content"]["google"]["thought_signature"] == "SIG" + + +def test_user_function_named_with_server_tool_arg_not_dropped(monkeypatch): + """Round 17: the OpenAI Responses translator must NOT drop a user function + whose JSON arguments contain `_server_tool: true` UNLESS the function name + is also a canonical builtin name. Otherwise a user schema with an + `_server_tool` field becomes invisible to the model.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "user_function", + "arguments": json.dumps({"_server_tool": True, "q": "x"}), + }, + } + ], + }, + { + "role": "tool", + "content": "result", + "tool_call_id": "call_user", + "name": "user_function", + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + fn_outs = [i for i in items if i.get("type") == "function_call_output"] + # User function call must survive (call + output). + assert any(c.get("name") == "user_function" for c in fn_calls), items + assert len(fn_outs) == 1, items + + +def test_builtin_named_with_server_tool_marker_dropped(monkeypatch): + """Round 17 control: a builtin (web_search) tagged with `_server_tool: + true` continues to be filtered from outbound history.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"_server_tool": True, "query": "x"}), + }, + } + ], + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + # Builtin server-side tool call must be filtered out. + assert all(c.get("name") != "web_search" for c in fn_calls), items + + +def test_gemini_tool_choice_none_disables_hosted_builtins(monkeypatch): + """Round 18: `tool_choice="none"` must drop hosted Google Search / code + execution from the Gemini body, not just user function declarations. + Otherwise an API client that opted out of tool use still triggers grounded + search (privacy + billing).""" + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search", "code_execution"], + tool_choice = "none", + ) + assert captured["body"].get("tools") is None, captured["body"] + + +def test_gemini_tool_choice_none_disables_function_declarations(monkeypatch): + """Round 18: `tool_choice="none"` must drop user function declarations as + well as hosted builtins from the Gemini body.""" + captured = _capture_body( + monkeypatch, + tool_choice = "none", + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + ) + assert captured["body"].get("tools") is None, captured["body"] + + +def test_schema_anyof_multitype_with_null_keeps_anyof_and_nullable(monkeypatch): + """Round 18: multi-branch unions with null (e.g. `Union[str, int, None]`) + must keep the slim anyOf without the null branch and add `nullable: true`; + Gemini rejects `{"type":"null"}` inside anyOf.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "either": { + "anyOf": [ + {"type": "string"}, + {"type": "integer"}, + {"type": "null"}, + ] + }, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + either = decls[0]["parameters"]["properties"]["either"] + assert either.get("nullable") is True + inner = either.get("anyOf") + assert isinstance(inner, list) and len(inner) == 2, either + assert all(not (isinstance(b, dict) and b.get("type") == "null") for b in inner), inner + + +def test_safe_fetch_image_redirect_malformed_url_no_crash(monkeypatch): + """Round 18: when the upstream 302 Location is a malformed bracketed-IPv6 + URL, the helper must return None instead of letting a urlparse ValueError + abort the chat stream.""" + import socket + import urllib.error + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubOpener: + def open( + self, + req, + timeout = None, + ): + raise urllib.error.HTTPError( + req.full_url, + 302, + "Found", + {"Location": "https://[bad/x.png"}, + None, + ) + + monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + + res = _drive(ep_mod._safe_fetch_image_for_gemini("https://cdn.example.com/x.png", "image/png")) + assert res is None + + +def test_safe_fetch_image_malformed_port_no_crash(): + """Round 18: a URL with a non-numeric port (`https://h:bad/x.png`) must + not raise; urlparse's port property lazily ValueErrors.""" + res = _drive(ep_mod._safe_fetch_image_for_gemini("https://example.com:bad/x.png", "image/png")) + assert res is None + + +def test_safe_fetch_image_missing_content_type_uses_fallback(monkeypatch): + """Round 18: when the server returns image bytes but no Content-Type + header, the helper must use the caller-provided fallback MIME (guessed from + URL extension) instead of dropping the image as `non-image + content-type=`.""" + import socket + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + headers = {"content-length": "3"} + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + return b"PNG" + + class _StubOpener: + def open( + self, + req, + timeout = None, + ): + return _StubResp() + + monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini("https://cdn.example.com/cat.png", "image/png") + ) + assert res is not None + assert res[0] == "image/png" + + +def test_anthropic_translates_openai_tool_calls_into_tool_use_blocks(monkeypatch): + """Round 18: an assistant turn with OpenAI-style top-level `tool_calls` + must be translated into Anthropic native `{type:"tool_use", id, name, + input}` content blocks before forwarding. The OpenAI `role="tool"` + follow-up must become a `role:"user"` message with a `tool_result` + block.""" + captured: dict = {"messages": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["messages"] = body.get("messages") + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "look up X"}, + { + "role": "assistant", + "content": "let me check", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"q":"x"}', + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_a", + "name": "lookup", + }, + {"role": "user", "content": "summarise"}, + ], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + msgs = captured["messages"] or [] + # No top-level tool_calls should remain. + assert all("tool_calls" not in m for m in msgs), msgs + # The assistant turn must now have content blocks including a tool_use + # block. + asst = [m for m in msgs if m.get("role") == "assistant"] + assert asst and isinstance(asst[0]["content"], list), asst + tool_uses = [b for b in asst[0]["content"] if b.get("type") == "tool_use"] + assert len(tool_uses) == 1, asst[0] + assert tool_uses[0]["name"] == "lookup" + assert tool_uses[0]["input"] == {"q": "x"} + # The role="tool" message must become a user/tool_result message. + tool_results: list[dict] = [] + for m in msgs: + if m.get("role") == "user" and isinstance(m.get("content"), list): + tool_results.extend(b for b in m["content"] if b.get("type") == "tool_result") + assert any( + tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text" + for tr in tool_results + ), msgs + + +def test_unmarked_user_web_search_function_survives_serialization(): + """Round 18: a user-defined function literally named `web_search` with NO + `_server_tool` marker must survive `_build_external_messages` when + forwarded to a non-native provider; only marked synthetic builtin cards may + be dropped.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"query": "x"}', + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert len(result) == 1, result + tcs = result[0].get("tool_calls") or [] + assert len(tcs) == 1, result + assert tcs[0]["function"]["name"] == "web_search" + + +def test_marked_server_builtin_dropped_from_build_external_messages(): + """Round 18: when a Gemini-native turn carrying a marked `image_generation` + server-tool card is forwarded to OpenAI / a custom Gemini OAI-compat proxy, + the tool_call must be dropped, not just have its extra_content stripped. + Forwarding an orphan `image_generation` tool_call would 400 the receiving + API.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "kind": "image"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "image_generation", + "arguments": marked_args, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + # Non-native providers: marked builtin tool_call must be dropped, and if it + # was the only payload, the whole message disappears. + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + # Empty assistant turn with only synthetic tool_call dropped. + assert result == [] or all(not (m.get("tool_calls") or []) for m in result), ( + provider_type, + result, + ) + + # Native Gemini preserves it (round-trips via extra_content). + result_native = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + assert len(result_native) == 1 + assert result_native[0]["tool_calls"][0]["function"]["name"] == "image_generation" + + +def test_openai_responses_tool_choice_none_drops_hosted_tools(monkeypatch): + """Round 18: `tool_choice="none"` must also drop hosted OpenAI Responses + builtins (web_search, code execution shell, image generation), not just + user function tools.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + enabled_tools = ["web_search", "code_execution", "image_generation"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("tools") in (None, []), body + + +def test_anthropic_tool_choice_none_drops_hosted_tools(monkeypatch): + """Round 19: tool_choice="none" must opt out of Anthropic hosted builtins + (web_search, web_fetch, code_execution) like it does for Gemini and OpenAI + Responses.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "web_fetch", "code_execution"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("tools") in (None, []), body + + +def test_openrouter_tool_choice_none_drops_web_plugin(monkeypatch): + """Round 19: tool_choice="none" must drop the OpenRouter web plugin so a + request that opted out of tool use doesn't still trigger hosted web + search.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("plugins") in (None, []), body + + +def test_kimi_tool_choice_none_skips_web_search_helper(monkeypatch): + """Round 19: when tool_choice="none" plus enabled_tools=["web_search"] on + Kimi, the dispatcher must NOT route into `_stream_kimi_web_search`. Falling + through to the generic OAI-compat path is expected.""" + routed_to_helper = {"called": False} + + real_helper = ExternalProviderClient._stream_kimi_web_search + + async def fake_helper(self, *args, **kwargs): # noqa: ARG001 + routed_to_helper["called"] = True + if False: + yield "" # pragma: no cover + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + fake_helper, + ) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "kimi", + base_url = "https://api.moonshot.ai/v1", + api_key = "sk-kimi-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "kimi-k2.6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + assert routed_to_helper["called"] is False + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + real_helper, + ) + + +def test_user_code_execution_function_not_dropped(): + """Round 19: a user-declared function literally named `code_execution` with + normal `code` arguments must survive `_build_external_messages` -- round + 17's shape heuristic dropped it, breaking function-calling round-trips.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "code_execution", + "arguments": '{"code": "print(1)"}', + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert len(result) == 1, result + tcs = result[0].get("tool_calls") or [] + assert len(tcs) == 1, result + assert tcs[0]["function"]["name"] == "code_execution" + + +def test_native_part_code_execution_treated_as_server_side(): + """Round 19: a Gemini `code_execution` card persists its replay payload at + `args.google.native_part` (no `_server_tool` marker on pre-PR cards). The + backend filter must still drop it for non-native providers because it's a + synthetic card, not a real user function.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + args_with_native_part = json.dumps( + { + "google": { + "native_part": { + "executableCode": { + "language": "PYTHON", + "code": "print(1)", + } + } + } + } + ) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": { + "name": "code_execution", + "arguments": args_with_native_part, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert result == [] or all(not (m.get("tool_calls") or []) for m in result), result + + +def test_remote_image_fetch_attempt_cap_includes_failures(monkeypatch): + """Round 19: the per-request image fetch count cap must count ATTEMPTS, + not just successes. Otherwise a request with 100 failing/slow URLs runs 100 + fetches each up to the 15s timeout.""" + fetch_calls: list[str] = [] + + async def fake_fetch( + url, + fallback_mime, + max_bytes = None, + ): + fetch_calls.append(url) + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + image_parts = [ + { + "type": "image_url", + "image_url": {"url": f"https://cdn.example.com/img{idx}.png"}, + } + for idx in range(20) + ] + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + *image_parts, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + assert len(fetch_calls) <= 8, len(fetch_calls) + + +def test_orphan_function_call_output_dropped_when_call_skipped(monkeypatch): + """Round 19: when a marked server-side builtin `function_call` is dropped + from OpenAI Responses input items, the matching role=tool follow-up must + also be dropped to avoid an orphan `function_call_output`.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"_server_tool": True, "query": "x"}), + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_b", + "name": "web_search", + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + fn_outs = [i for i in items if i.get("type") == "function_call_output"] + assert all(c.get("call_id") != "call_b" for c in fn_calls), items + assert all(o.get("call_id") != "call_b" for o in fn_outs), items + + +def test_schema_multitype_union_with_null_preserves_anyof(monkeypatch): + """Round 19: a JSON Schema `"type": ["string","integer","null"]` must be + sanitized to anyOf:[{string},{integer}] + nullable:true. Flattening to just + `{"type":"string"}` silently drops the integer branch and changes the + function contract.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "either": {"type": ["string", "integer", "null"]}, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + either = decls[0]["parameters"]["properties"]["either"] + assert either.get("nullable") is True + inner = either.get("anyOf") + assert isinstance(inner, list) and len(inner) == 2, either + types = sorted(b.get("type") for b in inner if isinstance(b, dict) and b.get("type")) + assert types == ["integer", "string"], inner + + +def test_invalid_gemini_model_rejected_before_image_fetch(monkeypatch): + """Round 19: invalid Gemini model IDs are rejected at the top of + `_stream_gemini`, BEFORE any user-controlled remote image fetch runs.""" + fetch_calls: list[str] = [] + + async def fake_fetch( + url, + fallback_mime, + max_bytes = None, + ): + fetch_calls.append(url) + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + { + "type": "image_url", + "image_url": {"url": "https://cdn.example.com/x.png"}, + }, + ], + } + ], + model = "../cachedContents/leak", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + assert fetch_calls == [], fetch_calls + + +def test_empty_assistant_turn_skipped_after_synthetic_tool_calls_dropped(): + """Round 20: when `_filter_tool_calls` drops every synthetic server-builtin + tool_call on an empty-content assistant turn, the whole message must be + skipped. Several providers reject `{"role":"assistant","content":""}` as an + empty assistant turn.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "kind": "image"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "image_generation", + "arguments": marked_args, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + # The empty assistant turn (only a synthetic builtin) must NOT appear + # in the output at all. + assert result == [], (provider_type, result) + + +def test_role_tool_dropped_when_matching_synthetic_call_filtered(): + """Round 20: `_build_external_messages` drops the matching role=tool + follow-up when its tool_call was a synthetic builtin that + `_filter_tool_calls` removed. Otherwise the receiving provider sees an + orphan tool_result with no tool_call.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "query": "x"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": marked_args, + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_b", + "name": "web_search", + }, + {"role": "user", "content": "continue"}, + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + # Only the user "continue" message survives. + roles = [m.get("role") for m in result] + assert roles == ["user"], result + + +def test_openrouter_no_synthetic_web_search_event_on_tool_choice_none(monkeypatch): + """Round 20: OpenRouter dispatcher must not emit synthetic web_search + tool_start / tool_end events when tool_choice="none"; otherwise the chat UI + shows a search card for a search that never happened.""" + captured_events: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :].strip() + if not payload or payload == "[DONE]": + continue + try: + obj = json.loads(payload) + except Exception: + continue + # Backend emits synthetic tool events as a top-level `_toolEvent` + # on the SSE payload (not nested inside `delta`). Read both shapes + # so a future format change can't mask this regression. + evt = obj.get("_toolEvent") + if isinstance(evt, dict): + captured_events.append(evt) + for ch in obj.get("choices") or []: + delta = ch.get("delta") or {} + nested = delta.get("_toolEvent") if isinstance(delta, dict) else None + if isinstance(nested, dict): + captured_events.append(nested) + await client.close() + + _drive(run()) + # No synthetic web_search tool_start / tool_end emitted. + assert all(e.get("tool_name") != "web_search" for e in captured_events), captured_events + + +def test_anthropic_role_tool_list_content_translates_to_tool_result(monkeypatch): + """Round 20: an OpenAI-shape role=tool message with list content + (`content=[{"type":"text","text":"result"}]`) must be translated into + Anthropic's native tool_result block, not forwarded as an invalid role=tool + message.""" + captured: dict = {"messages": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["messages"] = body.get("messages") + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "look up X"}, + { + "role": "assistant", + "content": "let me check", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"q":"x"}', + }, + } + ], + }, + { + "role": "tool", + "content": [{"type": "text", "text": "result_text"}], + "tool_call_id": "call_a", + "name": "lookup", + }, + {"role": "user", "content": "summarise"}, + ], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + msgs = captured["messages"] or [] + assert all(m.get("role") != "tool" for m in msgs), msgs + tool_results: list[dict] = [] + for m in msgs: + if m.get("role") == "user" and isinstance(m.get("content"), list): + tool_results.extend(b for b in m["content"] if b.get("type") == "tool_result") + assert any( + tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text" + for tr in tool_results + ), msgs + + +def test_data_url_non_image_mime_dropped(monkeypatch): + """Round 20: a `data:text/html;base64,...` image_url must be dropped from + the Gemini body, not forwarded as `inlineData.mimeType="text/html"` which + Gemini rejects.""" + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + { + "type": "image_url", + "image_url": { + "url": "data:text/html;base64,PGgxPmhpPC9oMT4=", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + assert not any("inlineData" in p for p in parts), parts + + +def test_youtube_filedata_uses_video_mime(monkeypatch): + """Round 20: YouTube `fileData.fileUri` must declare a video mimeType, not + `image/jpeg` guessed from the URL path.""" + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarise"}, + { + "type": "image_url", + "image_url": { + "url": "https://www.youtube.com/watch?v=abc", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + yt = next((p for p in parts if "fileData" in p), None) + assert yt is not None, parts + assert yt["fileData"]["mimeType"].startswith("video/"), yt + + +def test_openai_responses_assistant_text_serialized_before_function_call(monkeypatch): + """Round 20: in OpenAI Responses history, the assistant's visible text for + a turn that ALSO emitted a function_call must serialize BEFORE the + function_call item, matching the prior response.output sequence. Otherwise + function_call_output (the role=tool follow-up) appears to follow an + unrelated assistant message.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "Let me check that.", + "tool_calls": [ + { + "id": "call_w", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{}", + }, + } + ], + }, + { + "role": "tool", + "content": "sunny", + "tool_call_id": "call_w", + "name": "get_weather", + }, + {"role": "user", "content": "thanks"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + types = [i.get("type") or i.get("role") for i in items] + # Expected order: + # user ("weather?") + # assistant ("Let me check that.") + # function_call (get_weather) + # function_call_output (sunny) + # user ("thanks") + assert types == ["user", "assistant", "function_call", "function_call_output", "user"], items + + +def test_gemini_tool_choice_none_disables_image_generation(monkeypatch): + """Round 21: `tool_choice="none"` must also flip the implicit + image-generation hosted tool off on image-tier models. Otherwise + `responseModalities=["TEXT","IMAGE"]` still rides on the body and the + provider can generate (and bill for) image output despite the explicit + OpenAI tool opt-out.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tool_choice = "none", + ) + body = captured["body"] + assert body["generationConfig"].get("responseModalities") == ["TEXT"], body + + +def test_gemini_forced_function_tool_choice_drops_hosted_builtins(monkeypatch): + """Round 21: forced-function `tool_choice` (e.g. + `{"type":"function","function":{"name":"lookup"}}`) must suppress hosted + Google Search / code execution. Gemini's toolConfig only constrains + function declarations, not hosted tools, so leaving + `googleSearch`/`codeExecution` in `tools[]` lets them fire despite the + caller pinning a specific user function.""" + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search", "code_execution"], + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + tool_choice = { + "type": "function", + "function": {"name": "lookup"}, + }, + ) + body = captured["body"] + tool_kinds = [list(t.keys())[0] for t in (body.get("tools") or [])] + assert "googleSearch" not in tool_kinds, body + assert "codeExecution" not in tool_kinds, body + # User function declaration still survives. + assert "functionDeclarations" in tool_kinds, body + + +def test_gemini_forced_function_tool_choice_drops_image_generation(monkeypatch): + """Round 21: forced-function `tool_choice` must also flip the implicit + image-generation hosted tool off on image-tier models.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tool_choice = { + "type": "function", + "function": {"name": "lookup"}, + }, + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + ) + body = captured["body"] + assert body["generationConfig"].get("responseModalities") == ["TEXT"], body + + +def test_gemini_code_execution_native_part_list_replays_per_part_signatures(monkeypatch): + """Round 21: merged code-execution history must replay per-part + `thoughtSignature`s, not fan one top-level signature across every native + subpart. Gemini 3 strict validators reject a signature on the wrong + part.""" + history = [ + {"role": "user", "content": "plot 1+1"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(1+1)", + }, + "thoughtSignature": "SIG-EXEC", + }, + { + "codeExecutionResult": { + "id": "res_a", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + }, + ], + }, + }, + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + # Find the assistant turn replayed as native code-exec parts. + assistant_turn = next(c for c in contents if c["role"] == "model") + parts = assistant_turn["parts"] + exec_parts = [p for p in parts if "executableCode" in p] + result_parts = [p for p in parts if "codeExecutionResult" in p] + assert exec_parts and result_parts, parts + assert exec_parts[0].get("thoughtSignature") == "SIG-EXEC", exec_parts[0] + # codeExecutionResult had no signature -- must NOT inherit one. + assert "thoughtSignature" not in result_parts[0], result_parts[0] + + +def test_gemini_code_execution_legacy_merged_signature_only_on_executable(monkeypatch): + """Round 21: backward compat for pre-round-21 persisted history that stored + merged `native_part` as a single object plus a top-level + `thoughtSignature`. The replay branch must attach that signature only to + `executableCode` (where Gemini 3 emits it), not fan it across + `codeExecutionResult` / `inlineData`.""" + history = [ + {"role": "user", "content": "plot 1+1"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "executableCode": { + "id": "code_b", + "language": "PYTHON", + "code": "print(1+1)", + }, + "codeExecutionResult": { + "id": "res_b", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + "thoughtSignature": "LEGACY-SIG", + }, + }, + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_b", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + assistant_turn = next(c for c in contents if c["role"] == "model") + exec_parts = [p for p in assistant_turn["parts"] if "executableCode" in p] + result_parts = [p for p in assistant_turn["parts"] if "codeExecutionResult" in p] + assert exec_parts[0].get("thoughtSignature") == "LEGACY-SIG", exec_parts[0] + assert "thoughtSignature" not in result_parts[0], result_parts[0] + + +def test_gemini_role_tool_list_content_flattens_to_result_text(monkeypatch): + """Round 21: OpenAI-shape role=tool messages may carry list content like + `[{"type":"text","text":"result"}]`. Forwarding those parts verbatim into + `functionResponse.response.result` yields a list of content-part objects + instead of the actual tool output text.""" + history = [ + {"role": "user", "content": "look up"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "lookup", + "arguments": json.dumps({"q": "x"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "name": "lookup", + "content": [{"type": "text", "text": "answer-text"}], + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + fn_response = None + for c in contents: + for p in c.get("parts") or []: + if isinstance(p, dict) and "functionResponse" in p: + fn_response = p["functionResponse"] + break + if fn_response: + break + assert fn_response is not None, contents + assert fn_response["response"] == {"result": "answer-text"}, fn_response + + +def test_safe_fetch_image_threads_per_request_byte_budget(monkeypatch): + """Round 21: the aggregate per-request byte cap must be passed into + `_safe_fetch_image_for_gemini` so an oversize URL is refused via + Content-Length (short-circuit) rather than fully downloaded then + discarded.""" + import socket + + captured: dict = {"reads": 0, "content_length_seen": None} + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("8.8.8.8", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + # Declared 5 MiB, but caller passes a 1 MiB remaining budget. + headers = { + "content-type": "image/png", + "content-length": str(5 * 1024 * 1024), + } + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + captured["reads"] += 1 + return b"\x00" * (5 * 1024 * 1024) + + class _StubOpener: + def open( + self, + req, + timeout = None, + ): + return _StubResp() + + monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/big.png", + "image/png", + max_bytes = 1 * 1024 * 1024, + ) + ) + assert res is None + # Refused via Content-Length pre-check, never read. + assert captured["reads"] == 0 + + +def test_openai_chat_delta_type_includes_tool_calls_and_extra_content(): + """Round 21: the frontend `OpenAIChatDelta` interface must expose + `tool_calls` and `extra_content` so TypeScript callers can consume the + Gemini-native stream fields without `any` casts. A static-string assertion + against the .ts source; mirrors how other frontend wire-contract tests are + pinned from the backend suite.""" + import os + + here = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + types_path = os.path.join(here, "frontend", "src", "features", "chat", "types", "api.ts") + with open(types_path, "r", encoding = "utf-8") as f: + src = f.read() + assert "tool_calls?: OpenAIToolCallPart[]" in src, src[:200] + assert "extra_content?: Record" in src, src[:200] + assert "boolean | string | null" in src, src[:200] + + +def test_anthropic_forced_function_tool_choice_drops_hosted_tools(monkeypatch): + """Round 22: forced-function tool_choice must suppress Anthropic hosted + builtins like it does for Gemini. Pinning a user function + (`tool_choice={"type":"function","function":{"name":...}}`) while passing + `enabled_tools=["web_search","web_fetch","code_execution"]` should not still + fire those server-side.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "web_fetch", "code_execution"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + # No hosted tools in the body — only the caller's user-function + # declarations (none passed here). + tools = body.get("tools") or [] + hosted_tool_names = {"web_search", "web_fetch", "code_execution"} + for tool in tools: + assert tool.get("name") not in hosted_tool_names, body + + +def test_openrouter_forced_function_tool_choice_drops_web_plugin(monkeypatch): + """Round 22: forced-function tool_choice must drop the OpenRouter web + plugin too — caller pinned a user function, so OpenRouter must not attach + the hosted web-search plugin.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("plugins") in (None, []), body + + +def test_kimi_forced_function_tool_choice_skips_web_search_helper(monkeypatch): + """Round 22: forced-function tool_choice plus enabled_tools=["web_search"] + on Kimi must NOT route into `_stream_kimi_web_search`. Caller pinned a user + function; hosted $web_search should be suppressed for the same + privacy/billing reason.""" + routed_to_helper = {"called": False} + + async def fake_helper(self, *args, **kwargs): # noqa: ARG001 + routed_to_helper["called"] = True + if False: + yield "" # pragma: no cover + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + fake_helper, + ) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "kimi", + base_url = "https://api.moonshot.ai/v1", + api_key = "sk-kimi-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "kimi-k2.6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + assert not routed_to_helper["called"] + + +def test_openai_responses_forced_function_tool_choice_drops_hosted_tools(monkeypatch): + """Round 23: forced-function tool_choice on the OpenAI Responses path must + suppress hosted builtins (web_search, shell, image_generation) like it does + for Gemini / Anthropic / OpenRouter / Kimi. User-defined function tools + still flow through so the pinned function can resolve.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"event: response.completed\ndata: {}\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-openai-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "code_execution", "image_generation"], + tools = [ + { + "type": "function", + "function": { + "name": "lookup_record", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + tools = body.get("tools") or [] + hosted_types = {"web_search", "shell", "image_generation"} + hosted_seen = {t.get("type") for t in tools if isinstance(t, dict)} + assert not (hosted_seen & hosted_types), body + # The user function declaration must still be present so the pin has a + # target. + user_function_seen = any(isinstance(t, dict) and t.get("type") == "function" for t in tools) + assert user_function_seen, body + # And the forced-function tool_choice must be forwarded in Responses shape: + # `{type:"function", name:"..."}`. + tc = body.get("tool_choice") + assert isinstance(tc, dict) and tc.get("type") == "function", body + assert tc.get("name") == "lookup_record", body + + +def test_strip_provider_synthetic_tool_history_drops_text_only_extra_content(): + """Round 24: a plain text Gemini reply (no tool_calls) carrying + `extra_content.google.thought_signature` must still have that metadata + stripped before being forwarded to a local llama-server backend. Without + it, switching a Gemini thread mid-stream to a local GGUF model leaks + Gemini-only fields to llama-server.""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "Hello!", + "extra_content": {"google": {"thought_signature": "SIG_ABC"}}, + }, + {"role": "user", "content": "now in pirate voice"}, + ] + out = _strip_provider_synthetic_tool_history(messages) + # Same three turns, but the assistant's `extra_content` is gone. + assert [m["role"] for m in out] == ["user", "assistant", "user"] + assistant = out[1] + assert "extra_content" not in assistant, assistant + assert assistant["content"] == "Hello!" + + +def test_validate_and_resolve_host_blocks_shared_address_space(): + """Round 24 SSRF P1: 100.64.0.0/10 carrier-grade NAT addresses are + `is_private=False` AND `is_global=False` per Python's ipaddress docs. The + old denylist (is_private/loopback/link_local/etc.) missed them. Adding `not + ip.is_global` as the primary gate covers all non-public ranges, current and + future.""" + import socket as _socket + from core.inference import tools as _tools + + orig_getaddrinfo = _socket.getaddrinfo + + def fake_getaddrinfo(hostname, port, *args, **kwargs): + if hostname == "shared.example": + return [ + ( + _socket.AF_INET, + _socket.SOCK_STREAM, + 0, + "", + ("100.64.0.1", port), + ), + ] + return orig_getaddrinfo(hostname, port, *args, **kwargs) + + _socket.getaddrinfo = fake_getaddrinfo + try: + ok, reason, _ip = _tools._validate_and_resolve_host("shared.example", 443) + finally: + _socket.getaddrinfo = orig_getaddrinfo + assert ok is False, (ok, reason) + assert "non-public" in reason.lower() or "100.64.0.1" in reason + + +def test_gemini_custom_oai_compat_base_skips_native_allowlist(): + """Round 24: a custom Gemini OAI-compatible base (LiteLLM/proxy) must NOT + have its model list filtered through the native Gemini allowlist regex. A + LiteLLM gateway returning + `["google/gemini-2.5-flash", "my-team/gemini", "gemini-2.5-flash"]` should + pass through; the native filter would strip the prefixed IDs even though + chat dispatch routes them via the OpenAI-compatible client.""" + import asyncio as _asyncio + + from routes import providers as _providers + from routes.providers import ( + ProviderModelsRequest, + list_provider_models, + ) + + captured: dict = {"base": None} + + class _FakeClient: + def __init__(self, *, base_url, **kwargs): + captured["base"] = base_url + + async def list_models(self): + return [ + {"id": "google/gemini-2.5-flash"}, + {"id": "my-team/gemini"}, + {"id": "gemini-2.5-flash"}, + ] + + async def close(self): + return None + + orig = _providers.ExternalProviderClient + _providers.ExternalProviderClient = _FakeClient + try: + req = ProviderModelsRequest( + provider_type = "gemini", + base_url = "https://litellm.example/v1", + ) + result = _asyncio.run(list_provider_models(req, current_subject = "unsloth")) + finally: + _providers.ExternalProviderClient = orig + ids = {m.id for m in result} + # All three IDs survive — native allowlist bypassed. + assert "google/gemini-2.5-flash" in ids, ids + assert "my-team/gemini" in ids, ids + assert "gemini-2.5-flash" in ids, ids + + +def test_strip_provider_synthetic_tool_history_drops_synthetic_only(): + """Round 22: switching a thread from native Gemini (code_execution / + image_generation tool_cards in history) to a local GGUF backend must strip + the synthetic tool_calls + matching role=tool replies before llama-server + sees them. Real user-function tool_calls and their matching tool replies + must survive.""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "let me run it", + "tool_calls": [ + { + "id": "synth_ce_1", + "type": "function", + "function": { + "name": "code_execution", + "arguments": json.dumps( + { + "_server_tool": True, + "google": {"native_part": {"parts": []}}, + } + ), + }, + "extra_content": {"google": {"thought_signature": "abc"}}, + }, + { + "id": "real_lookup", + "type": "function", + "function": { + "name": "lookup_user", + "arguments": json.dumps({"id": 42}), + }, + }, + ], + "extra_content": {"google": {"thought_signature": "msglevel"}}, + }, + { + "role": "tool", + "tool_call_id": "synth_ce_1", + "content": "Gemini-only result text", + }, + { + "role": "tool", + "tool_call_id": "real_lookup", + "content": '{"name": "alice"}', + }, + ] + out = _strip_provider_synthetic_tool_history(messages) + assistant = next(m for m in out if m.get("role") == "assistant") + tcs = assistant["tool_calls"] + assert len(tcs) == 1, tcs + assert tcs[0]["id"] == "real_lookup" + assert "extra_content" not in tcs[0] + assert "extra_content" not in assistant + tool_msgs = [m for m in out if m.get("role") == "tool"] + assert len(tool_msgs) == 1 + assert tool_msgs[0]["tool_call_id"] == "real_lookup" + + +def test_strip_provider_synthetic_tool_history_drops_empty_assistant(): + """If every tool_call was synthetic and the assistant turn had no content, + the entire turn must be dropped (llama-server rejects empty assistant + messages with no tool_calls).""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "draw a sloth"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "synth_imggen", + "type": "function", + "function": { + "name": "image_generation", + "arguments": json.dumps( + { + "google": { + "native_part": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "Zm9v", + } + } + ] + } + } + } + ), + }, + } + ], + }, + {"role": "tool", "tool_call_id": "synth_imggen", "content": "(image)"}, + {"role": "user", "content": "now try in pirate voice"}, + ] + out = _strip_provider_synthetic_tool_history(messages) + roles = [m.get("role") for m in out] + # Synthetic assistant + its tool reply are both gone; only the two user + # turns survive. + assert roles == ["user", "user"], out + + +def test_openrouter_no_synthetic_web_search_event_on_forced_function_tool_choice(monkeypatch): + """Round 22 sibling of the round-20 `tool_choice='none'` test: when the + caller forces a specific function via `tool_choice={"type":"function", ...}` + AND passes `enabled_tools=["web_search"]`, the OpenRouter path must NOT + synthesize a fake `web_search` tool card. The plugin wasn't attached + upstream, so the UI must not see a server-tool card.""" + captured_events: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = (b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n' b"data: [DONE]\n\n"), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + payload = line.strip().removeprefix("data: ") + if payload and payload != "[DONE]": + try: + captured_events.append(json.loads(payload)) + except Exception: + pass + await client.close() + + _drive(run()) + for evt in captured_events: + for choice in evt.get("choices") or []: + delta = choice.get("delta") or {} + extra = delta.get("extra_content") or {} + tool_event = extra.get("toolEvent") if isinstance(extra, dict) else None + if isinstance(tool_event, dict): + assert tool_event.get("tool_name") != "web_search", evt diff --git a/studio/backend/tests/test_gguf_completion_usage.py b/studio/backend/tests/test_gguf_completion_usage.py new file mode 100644 index 0000000000..d1e05f3e0e --- /dev/null +++ b/studio/backend/tests/test_gguf_completion_usage.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression tests for GGUF non-streaming chat completion usage.""" + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from auth.authentication import get_current_subject +import routes.inference as inference_route + + +class _GgufBackend: + is_loaded = True + model_identifier = "test/model.gguf" + _is_audio = False + is_vision = False + supports_tools = False + + def __init__(self, usage): + self.usage = usage + + def generate_chat_completion(self, **kwargs): + yield "answer" + yield { + "type": "metadata", + "usage": self.usage, + "timings": {"prompt_n": 23, "predicted_n": 1283}, + } + + +def _request_completion(monkeypatch, usage): + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _GgufBackend(usage)) + monkeypatch.setattr(inference_route, "_effective_enable_tools", lambda payload: False) + + app = FastAPI() + app.include_router(inference_route.router) + app.dependency_overrides[get_current_subject] = lambda: "test-user" + + return TestClient(app).post( + "/chat/completions", + json = { + "messages": [{"role": "user", "content": "Why is the sky blue?"}], + "stream": False, + }, + ) + + +def test_non_streaming_gguf_completion_includes_generated_usage(monkeypatch): + response = _request_completion( + monkeypatch, + {"prompt_tokens": 23, "completion_tokens": 1283, "total_tokens": 1306}, + ) + + assert response.status_code == 200 + usage = response.json()["usage"] + assert usage["prompt_tokens"] == 23 + assert usage["completion_tokens"] == 1283 + assert usage["total_tokens"] == 1306 + assert usage["prompt_tokens_details"] == {"cached_tokens": 0, "audio_tokens": 0} + assert usage["completion_tokens_details"] == { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + } + + +def test_non_streaming_gguf_completion_defaults_nullable_usage_to_zero(monkeypatch): + response = _request_completion( + monkeypatch, + {"prompt_tokens": None, "completion_tokens": 1283, "total_tokens": None}, + ) + + assert response.status_code == 200 + usage = response.json()["usage"] + assert usage["prompt_tokens"] == 0 + assert usage["completion_tokens"] == 1283 + assert usage["total_tokens"] == 1283 + assert usage["prompt_tokens_details"] == {"cached_tokens": 0, "audio_tokens": 0} + assert usage["completion_tokens_details"] == { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + } diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index cf1a17347f..d3d4387720 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -1,8 +1,8 @@ # 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.""" +"""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 @@ -14,6 +14,7 @@ from utils.models.gguf_metadata import ( is_mmproj_by_metadata, pairing_score, read_gguf_general_metadata, + read_mmproj_audio_capability, ) @@ -21,6 +22,7 @@ _GGUF_MAGIC = 0x46554747 _VTYPE_STRING = 8 _VTYPE_UINT32 = 4 _VTYPE_ARRAY = 9 +_VTYPE_BOOL = 7 def _enc_string(s: str) -> bytes: @@ -33,9 +35,11 @@ def _enc_kv_string(key: str, value: str) -> bytes: def _enc_kv_uint32(key: str, value: int) -> bytes: - return ( - _enc_string(key) + struct.pack(" bytes: + return _enc_string(key) + struct.pack(" bytes: @@ -53,11 +57,15 @@ def _write_synthetic_gguf( *, extra_uint32: Mapping[str, int] | None = None, extra_string_arrays: Mapping[str, Iterable[str]] | None = None, + extra_bools: Mapping[str, bool] | 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) + extra_bools = extra_bools or {} + kv_count = ( + len(general_strings) + len(extra_uint32) + len(extra_string_arrays) + len(extra_bools) + ) body = b"" for k, v in general_strings.items(): body += _enc_kv_string(k, v) @@ -65,6 +73,8 @@ def _write_synthetic_gguf( body += _enc_kv_uint32(k, v) for k, v in extra_string_arrays.items(): body += _enc_kv_string_array(k, v) + for k, v in extra_bools.items(): + body += _enc_kv_bool(k, v) header = struct.pack( " str: + return "data: " + json.dumps({"choices": [{"index": 0, "delta": delta}]}) + "\n" + + +def _done() -> str: + return "data: [DONE]\n" + + +def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._process = object() + backend._healthy = True + backend._port = 48848 + backend._api_key = None + backend._effective_context_length = 4096 + backend._supports_reasoning = False + backend._reasoning_always_on = False + backend._reasoning_style = "enable_thinking" + backend._supports_preserve_thinking = False + + @contextlib.contextmanager + def fake_stream_with_retry( + _client, + _url, + payload, + _cancel_event, + headers = None, + ): + payloads.append(copy.deepcopy(payload)) + yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() + + def fake_iter_text_cancellable(response, _cancel_event): + yield from response.chunks + + monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) + monkeypatch.setattr(backend, "_iter_text_cancellable", fake_iter_text_cancellable) + return backend + + +def _replay_route_cursor(events: list[dict]) -> dict: + """Replicate the GGUF SSE route's cumulative-cursor loop. + + Mirrors studio/backend/routes/inference.py: reset ``prev_text`` on empty + status and on ``tool_start``; otherwise diff each cumulative ``content`` + snapshot against the cursor and stream the delta. The preface/final text + here carry no tool XML, so the display strip is the identity -- the cursor + reset is the behaviour under test. + """ + prev_text = "" + visible_deltas: list[str] = [] + tool_starts: list[dict] = [] + statuses: list[str] = [] + for event in events: + etype = event["type"] + if etype == "status": + if not event["text"]: + prev_text = "" + statuses.append(event["text"]) + continue + if etype in ("tool_start", "tool_end"): + if etype == "tool_start": + prev_text = "" + tool_starts.append(event) + continue + if etype == "metadata": + continue + clean_cumulative = event.get("text", "") + new_text = clean_cumulative[len(prev_text) :] + prev_text = clean_cumulative + if not new_text: + continue + visible_deltas.append(new_text) + return { + "visible": "".join(visible_deltas), + "tool_starts": tool_starts, + "statuses": statuses, + } + + +def _replay_route_cursor_without_status_reset(events: list[dict]) -> dict: + """Pre-fix control: identical to the route loop but never resets the + cursor on an empty status (only on ``tool_start``).""" + prev_text = "" + visible_deltas: list[str] = [] + for event in events: + etype = event["type"] + if etype == "status": + continue + if etype in ("tool_start", "tool_end"): + if etype == "tool_start": + prev_text = "" + continue + if etype == "metadata": + continue + clean_cumulative = event.get("text", "") + new_text = clean_cumulative[len(prev_text) :] + prev_text = clean_cumulative + if not new_text: + continue + visible_deltas.append(new_text) + return {"visible": "".join(visible_deltas)} + + +def _web_search_tool() -> dict: + return { + "type": "function", + "function": { + "name": "web_search", + "description": "Search the web.", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, + } + + +def test_final_answer_survives_preface_then_disabled_tool_noop(monkeypatch): + """Preface text, then a call to a disabled tool (internal no-op). + + A disabled-tool decision emits no ``tool_start`` and forces the final + no-tools pass. The route cursor must be reset before that pass so the + short final answer is not diffed away against the longer preface. + """ + preface = "Let me run a quick command to double-check." + final = "All set." # deliberately shorter than the preface -> truncation is visible + + # Single turn: visible preface + a call to `terminal`, which is NOT in the + # enabled tool list, so the controller marks it disabled -> internal no-op. + turn_stream = [ + _sse({"content": preface}), + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_disabled", + "type": "function", + "function": { + "name": "terminal", + "arguments": json.dumps({"command": "ls"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": final}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [turn_stream, final_stream], payloads) + + executed: list[str] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_kw: executed.append(name) or "should-not-run", + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "answer me"}], + tools = [_web_search_tool()], # terminal intentionally absent + temperature = 0.0, + max_tool_iterations = 5, + ) + ) + + replay = _replay_route_cursor(events) + + # Disabled tool is an internal no-op: never executed, no visible card. + assert executed == [] + assert replay["tool_starts"] == [] + + # The generator must emit an empty status that resets the route cursor + # before the final pass; otherwise `final` (shorter than `preface`) would + # be diffed to nothing and dropped. + assert "" in replay["statuses"], "no cursor-resetting empty status emitted" + + # Both the preface and the final answer survive, in order, untruncated. + assert preface in replay["visible"], replay["visible"] + assert final in replay["visible"], replay["visible"] + assert replay["visible"].index(preface) < replay["visible"].index(final) + assert replay["visible"].count(preface) == 1 + + # Negative control: a route loop that does NOT reset on empty status (the + # pre-fix behaviour) would diff `final` against the stale preface cursor + # and drop it -- proving the empty status is load-bearing here. + no_reset = _replay_route_cursor_without_status_reset(events) + assert final not in no_reset["visible"], no_reset["visible"] diff --git a/studio/backend/tests/test_gguf_routing.py b/studio/backend/tests/test_gguf_routing.py new file mode 100644 index 0000000000..8a8bff69aa --- /dev/null +++ b/studio/backend/tests/test_gguf_routing.py @@ -0,0 +1,100 @@ +# 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 GGUF routing in detect_gguf_model. + +Regression test: on Windows a .gguf file can briefly appear inaccessible +during llama-server teardown, making is_file() return False and routing +the model to the transformers backend instead of llama-server. +""" + +import sys +import os +import types +from pathlib import Path +from unittest.mock import patch + +# Stub structlog before importing backend modules (as in other suite tests) +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _): + return lambda *a, **k: None + + sys.modules["structlog"] = types.SimpleNamespace( + get_logger = lambda *a, **k: _DummyLogger(), + BoundLogger = _DummyLogger, + ) + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from utils.models.model_config import detect_gguf_model + + +def test_detects_gguf_file_normally(tmp_path): + """Normal case: .gguf file exists and is accessible.""" + gguf = tmp_path / "gpt-oss-20b-MXFP4.gguf" + gguf.write_bytes(b"") + result = detect_gguf_model(str(gguf)) + assert result is not None + assert result.endswith("gpt-oss-20b-MXFP4.gguf") + + +def test_detects_gguf_when_stat_raises_oserror(tmp_path): + """ + Regression: on Windows is_file()/exists() call stat(), which raises OSError + in the brief lock window after llama-server is killed. detect_gguf_model must + still route to llama-server by file extension alone. + """ + gguf = tmp_path / "gpt-oss-20b-MXFP4.gguf" + gguf.write_bytes(b"") + + original_stat = Path.stat + + def flaky_stat(self, **kwargs): + if self.suffix.lower() == ".gguf": + raise OSError("file temporarily inaccessible (Windows lock window)") + return original_stat(self, **kwargs) + + with patch.object(Path, "stat", flaky_stat): + result = detect_gguf_model(str(gguf)) + + assert result is not None, ( + "detect_gguf_model returned None when stat() raised OSError. " + "This causes the model to fall through to the transformers backend." + ) + + +def test_does_not_detect_mmproj_as_main_model(tmp_path): + """mmproj files must never be returned as the primary model.""" + mmproj = tmp_path / "mmproj-model-f16.gguf" + mmproj.write_bytes(b"") + result = detect_gguf_model(str(mmproj)) + assert result is None + + +def test_detects_gguf_in_directory(tmp_path): + """Directory containing a .gguf file is resolved to that file.""" + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"") + result = detect_gguf_model(str(tmp_path)) + assert result is not None + assert result.endswith("model-Q4_K_M.gguf") + + +def test_directory_named_like_gguf_scans_inside(tmp_path): + """A directory named *.gguf resolves the real .gguf inside, not itself.""" + gguf_dir = tmp_path / "mymodel.gguf" + gguf_dir.mkdir() + inner = gguf_dir / "model-Q4_K_M.gguf" + inner.write_bytes(b"") + result = detect_gguf_model(str(gguf_dir)) + assert result is not None + assert result.endswith("model-Q4_K_M.gguf") + + +def test_returns_none_for_non_gguf_path(tmp_path): + """Non-.gguf paths with no .gguf files inside return None.""" + result = detect_gguf_model(str(tmp_path)) + assert result is None diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index a1fe5653ef..cb26330ed0 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -64,9 +64,7 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): def test_parent_visibility_uses_empty_numeric_ids_for_uuid_masks(self): with ( - patch.dict( - os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True - ), + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): self.assertEqual(get_parent_visible_gpu_ids(), []) @@ -96,9 +94,7 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): def test_explicit_ids_are_rejected_for_uuid_parent_visibility(self): with ( - patch.dict( - os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True - ), + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): with self.assertRaisesRegex( @@ -185,7 +181,7 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase): self.assertAlmostEqual(result["devices"][1]["memory_total_gb"], 29.3, places = 1) def test_uuid_parent_visibility_falls_back_to_torch(self): - """UUID/MIG masks should fall through nvidia to torch fallback and + """UUID/MIG masks fall through nvidia to the torch fallback and still report visible devices using relative ordinals.""" fake_torch_devices = [ { @@ -204,13 +200,9 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase): }, ] with ( - patch.dict( - os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True - ), + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), - patch( - "utils.hardware.hardware._torch_get_physical_gpu_count", return_value = 2 - ), + patch("utils.hardware.hardware._torch_get_physical_gpu_count", return_value = 2), patch( "utils.hardware.hardware._torch_get_per_device_info", return_value = fake_torch_devices, @@ -254,9 +246,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_get_device_map_uses_all_inherited_visible_gpus_for_uuid_masks(self): with ( - patch.dict( - os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True - ), + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), ): self.assertEqual(get_device_map(None), "balanced") @@ -376,9 +366,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): return_value = 1234, ), ): - model_size_bytes, source = _hw_module.estimate_fp16_model_size_bytes( - "unsloth/test" - ) + model_size_bytes, source = _hw_module.estimate_fp16_model_size_bytes("unsloth/test") self.assertEqual(model_size_bytes, 1234) self.assertEqual(source, "vllm_utils") @@ -475,9 +463,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self): with ( - patch.dict( - os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True - ), + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", return_value = ( @@ -524,9 +510,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): patch( "core.training.training._CTX.Process", return_value = DummyProcess() ) as mock_process, - patch( - "core.training.training.threading.Thread", return_value = DummyThread() - ), + patch("core.training.training.threading.Thread", return_value = DummyThread()), ): backend.start_training( job_id = "test-job-1", @@ -567,9 +551,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): patch( "core.training.training._CTX.Process", return_value = DummyProcess() ) as mock_process, - patch( - "core.training.training.threading.Thread", return_value = DummyThread() - ), + patch("core.training.training.threading.Thread", return_value = DummyThread()), ): backend.start_training( job_id = "test-job-2", @@ -599,9 +581,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): dummy_queue = object() with ( - patch.dict( - os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True - ), + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), patch( "core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue], @@ -609,9 +589,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): patch( "core.training.training._CTX.Process", return_value = DummyProcess() ) as mock_process, - patch( - "core.training.training.threading.Thread", return_value = DummyThread() - ), + patch("core.training.training.threading.Thread", return_value = DummyThread()), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", return_value = ( @@ -629,9 +607,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): config = mock_process.call_args.kwargs["kwargs"]["config"] self.assertIsNone(config["resolved_gpu_ids"]) - self.assertEqual( - config["gpu_selection"]["selection_mode"], "inherit_parent_visible" - ) + self.assertEqual(config["gpu_selection"]["selection_mode"], "inherit_parent_visible") def test_inference_orchestrator_resolves_explicit_gpu_ids_before_spawn(self): class DummyThread: @@ -643,7 +619,6 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): with patch("core.inference.orchestrator.threading.Thread", DummyThread): from core.inference.orchestrator import InferenceOrchestrator - orchestrator = InferenceOrchestrator() config = SimpleNamespace(identifier = "unsloth/test", gguf_variant = None) @@ -660,9 +635,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): "_wait_response", return_value = {"success": True, "model_info": {}}, ), - patch( - "utils.transformers_version.needs_transformers_5", return_value = False - ), + patch("utils.transformers_version.needs_transformers_5", return_value = False), ): self.assertTrue(orchestrator.load_model(config = config, gpu_ids = [1])) @@ -681,7 +654,6 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): with patch("core.inference.orchestrator.threading.Thread", DummyThread): from core.inference.orchestrator import InferenceOrchestrator - orchestrator = InferenceOrchestrator() config = SimpleNamespace(identifier = "unsloth/test", gguf_variant = None) @@ -698,9 +670,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): "_wait_response", return_value = {"success": True, "model_info": {}}, ), - patch( - "utils.transformers_version.needs_transformers_5", return_value = False - ), + patch("utils.transformers_version.needs_transformers_5", return_value = False), ): self.assertTrue(orchestrator.load_model(config = config, gpu_ids = None)) @@ -782,9 +752,7 @@ class TestRouteErrors(unittest.TestCase): raise ValueError("Invalid gpu_ids [99]") with ( - patch.object( - training_route, "get_training_backend", return_value = DummyBackend() - ), + patch.object(training_route, "get_training_backend", return_value = DummyBackend()), patch( "core.inference.get_inference_backend", return_value = SimpleNamespace(active_model_name = None), @@ -795,9 +763,7 @@ class TestRouteErrors(unittest.TestCase): ), ): with self.assertRaises(HTTPException) as exc_info: - asyncio.run( - training_route.start_training(request, current_subject = "test-user") - ) + asyncio.run(training_route.start_training(request, current_subject = "test-user")) self.assertEqual(exc_info.exception.status_code, 400) self.assertIn("gpu_ids [99]", exc_info.exception.detail) @@ -826,9 +792,7 @@ class TestRouteErrors(unittest.TestCase): ) with ( - patch.object( - training_route, "get_training_backend", return_value = DummyBackend() - ), + patch.object(training_route, "get_training_backend", return_value = DummyBackend()), patch( "core.inference.get_inference_backend", return_value = SimpleNamespace(active_model_name = None), @@ -839,9 +803,7 @@ class TestRouteErrors(unittest.TestCase): ), ): with self.assertRaises(HTTPException) as exc_info: - asyncio.run( - training_route.start_training(request, current_subject = "test-user") - ) + asyncio.run(training_route.start_training(request, current_subject = "test-user")) self.assertEqual(exc_info.exception.status_code, 400) self.assertIn("UUID/MIG", exc_info.exception.detail) @@ -976,22 +938,17 @@ class TestRouteErrors(unittest.TestCase): class TestRaiseIfOffloaded(unittest.TestCase): def test_no_offload_is_noop(self): from utils.hardware import raise_if_offloaded - model = SimpleNamespace(hf_device_map = {"model.embed_tokens": 0, "lm_head": 1}) raise_if_offloaded(model, "balanced", "Test") def test_cpu_offload_raises(self): from utils.hardware import raise_if_offloaded - - model = SimpleNamespace( - hf_device_map = {"model.layers.0": 0, "model.layers.1": "cpu"} - ) + model = SimpleNamespace(hf_device_map = {"model.layers.0": 0, "model.layers.1": "cpu"}) with self.assertRaisesRegex(ValueError, "offloaded"): raise_if_offloaded(model, "balanced", "Test") def test_no_device_map_attr_is_noop(self): from utils.hardware import raise_if_offloaded - raise_if_offloaded(SimpleNamespace(), "sequential", "Test") @@ -1314,7 +1271,6 @@ class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase): config = object(), ): from utils.hardware import hardware as hardware_module - with ( patch.object( hardware_module, @@ -1382,7 +1338,7 @@ class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase): self.assertEqual(src, "weight_bytes") def test_equal_local_and_config_keeps_config_label(self): - # why: tie-breaker is "local must be strictly larger" so an exact + # Tie-breaker is "local must be strictly larger", so an exact # match keeps the config-derived path. same = 8 * (1 << 30) bytes_, src = self._run( @@ -1395,7 +1351,6 @@ class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase): def test_remote_safetensors_path_unaffected_by_local_weights(self): from utils.hardware import hardware as hardware_module - with ( patch.object( hardware_module, diff --git a/studio/backend/tests/test_gpu_selection_sandbox.py b/studio/backend/tests/test_gpu_selection_sandbox.py index 830a98a2fb..733933271b 100644 --- a/studio/backend/tests/test_gpu_selection_sandbox.py +++ b/studio/backend/tests/test_gpu_selection_sandbox.py @@ -2,9 +2,9 @@ """ Sandbox test for multi-GPU selection logic. -Tests the core GPU selection, memory estimation, and device_map logic -in an isolated environment. Can be run on Linux, macOS, and Windows -without requiring actual GPUs -- all hardware calls are mocked. +Tests GPU selection, memory estimation, and device_map logic in +isolation. Runs on Linux, macOS, and Windows without real GPUs -- all +hardware calls are mocked. Usage: python -m pytest studio/backend/tests/test_gpu_selection_sandbox.py -v @@ -18,7 +18,7 @@ import unittest from pathlib import Path from unittest.mock import patch, MagicMock -# Ensure backend is on sys.path +# Ensure backend is on sys.path. _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) @@ -33,9 +33,8 @@ def _make_fake_config( num_key_value_heads = 8, tie_word_embeddings = False, ): - """Create a fake HF config-like object for estimation tests.""" + """Fake HF config-like object for estimation tests.""" from types import SimpleNamespace - return SimpleNamespace( vocab_size = vocab_size, hidden_size = hidden_size, @@ -48,7 +47,7 @@ def _make_fake_config( class TestEstimateFP16ModelSizeFromConfig(unittest.TestCase): - """Test the config-based model size estimation.""" + """Config-based model size estimation.""" def test_llama_8b_size_reasonable(self): from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config @@ -91,7 +90,7 @@ class TestEstimateFP16ModelSizeFromConfig(unittest.TestCase): from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config from types import SimpleNamespace - config = SimpleNamespace(vocab_size = 32000) # Missing most fields + config = SimpleNamespace(vocab_size = 32000) # most fields missing size = _estimate_fp16_model_size_bytes_from_config(config) self.assertIsNone(size) @@ -118,11 +117,10 @@ class TestEstimateFP16ModelSizeFromConfig(unittest.TestCase): class TestEstimateRequiredModelMemory(unittest.TestCase): - """Test memory requirement estimation.""" + """Memory requirement estimation.""" def test_inference_fp16_uses_1_3x(self): from utils.hardware.hardware import estimate_required_model_memory_gb - with patch( "utils.hardware.hardware.estimate_fp16_model_size_bytes", return_value = (10 * (1024**3), "config"), # 10GB model @@ -138,7 +136,6 @@ class TestEstimateRequiredModelMemory(unittest.TestCase): def test_inference_4bit_uses_reduced_estimate(self): from utils.hardware.hardware import estimate_required_model_memory_gb - with patch( "utils.hardware.hardware.estimate_fp16_model_size_bytes", return_value = (30 * (1024**3), "config"), # 30GB fp16 model @@ -154,7 +151,6 @@ class TestEstimateRequiredModelMemory(unittest.TestCase): def test_4bit_training_reduces_base(self): from utils.hardware.hardware import estimate_required_model_memory_gb - with patch( "utils.hardware.hardware.estimate_fp16_model_size_bytes", return_value = (30 * (1024**3), "config"), # 30GB fp16 model @@ -170,7 +166,6 @@ class TestEstimateRequiredModelMemory(unittest.TestCase): def test_full_finetune_uses_3_5x(self): from utils.hardware.hardware import estimate_required_model_memory_gb - with patch( "utils.hardware.hardware.estimate_fp16_model_size_bytes", return_value = (10 * (1024**3), "config"), # 10GB model @@ -185,7 +180,6 @@ class TestEstimateRequiredModelMemory(unittest.TestCase): def test_returns_none_when_unavailable(self): from utils.hardware.hardware import estimate_required_model_memory_gb - with patch( "utils.hardware.hardware.estimate_fp16_model_size_bytes", return_value = (None, "unavailable"), @@ -195,10 +189,10 @@ class TestEstimateRequiredModelMemory(unittest.TestCase): class TestAutoSelectGpuIds(unittest.TestCase): - """Test automatic GPU selection based on model size and free memory.""" + """Automatic GPU selection by model size and free memory.""" def _make_utilization(self, devices): - """Create a fake utilization response.""" + """Fake utilization response.""" return { "available": True, "devices": [ @@ -214,7 +208,6 @@ class TestAutoSelectGpuIds(unittest.TestCase): def test_single_gpu_sufficient(self): from utils.hardware.hardware import auto_select_gpu_ids import utils.hardware.hardware as hw - with ( patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), patch.object( @@ -261,7 +254,6 @@ class TestAutoSelectGpuIds(unittest.TestCase): def test_two_gpus_needed(self): from utils.hardware.hardware import auto_select_gpu_ids import utils.hardware.hardware as hw - with ( patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), patch.object( @@ -305,7 +297,6 @@ class TestAutoSelectGpuIds(unittest.TestCase): def test_non_cuda_returns_none(self): from utils.hardware.hardware import auto_select_gpu_ids import utils.hardware.hardware as hw - with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): selected, meta = auto_select_gpu_ids("test/model") self.assertIsNone(selected) @@ -313,12 +304,11 @@ class TestAutoSelectGpuIds(unittest.TestCase): class TestGetDeviceMap(unittest.TestCase): - """Test device_map string generation.""" + """device_map string generation.""" def test_single_gpu_returns_sequential(self): from utils.hardware.hardware import get_device_map import utils.hardware.hardware as hw - with ( patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), patch.object( @@ -338,7 +328,6 @@ class TestGetDeviceMap(unittest.TestCase): def test_multi_gpu_returns_balanced(self): from utils.hardware.hardware import get_device_map import utils.hardware.hardware as hw - with patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA): dm = get_device_map(gpu_ids = [0, 1]) self.assertEqual(dm, "balanced") @@ -346,18 +335,16 @@ class TestGetDeviceMap(unittest.TestCase): def test_cpu_returns_sequential(self): from utils.hardware.hardware import get_device_map import utils.hardware.hardware as hw - with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): dm = get_device_map(gpu_ids = None) self.assertEqual(dm, "sequential") class TestResolveRequestedGpuIds(unittest.TestCase): - """Test GPU ID validation.""" + """GPU ID validation.""" def test_none_returns_parent_visible(self): from utils.hardware.hardware import resolve_requested_gpu_ids - with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), @@ -367,7 +354,6 @@ class TestResolveRequestedGpuIds(unittest.TestCase): def test_empty_list_returns_parent_visible(self): from utils.hardware.hardware import resolve_requested_gpu_ids - with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), @@ -377,7 +363,6 @@ class TestResolveRequestedGpuIds(unittest.TestCase): def test_duplicates_rejected(self): from utils.hardware.hardware import resolve_requested_gpu_ids - with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1,2"}, clear = False), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), @@ -387,7 +372,6 @@ class TestResolveRequestedGpuIds(unittest.TestCase): def test_out_of_range_rejected(self): from utils.hardware.hardware import resolve_requested_gpu_ids - with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"}, clear = False), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4), @@ -397,11 +381,8 @@ class TestResolveRequestedGpuIds(unittest.TestCase): def test_uuid_env_var_rejects_explicit_ids(self): from utils.hardware.hardware import resolve_requested_gpu_ids - with ( - patch.dict( - os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False - ), + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): with self.assertRaises(ValueError): @@ -409,11 +390,10 @@ class TestResolveRequestedGpuIds(unittest.TestCase): class TestApplyGpuIds(unittest.TestCase): - """Test CUDA_VISIBLE_DEVICES environment variable setting.""" + """CUDA_VISIBLE_DEVICES environment variable setting.""" def test_apply_list(self): from utils.hardware.hardware import apply_gpu_ids - with patch.dict(os.environ, {}, clear = False): apply_gpu_ids([3, 5]) self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), "3,5") @@ -427,10 +407,10 @@ class TestApplyGpuIds(unittest.TestCase): class TestMultiGpuOverheadAccounting(unittest.TestCase): - """Test that multi-GPU overhead is applied correctly. + """Multi-GPU overhead is applied correctly. - The first GPU should keep its full free memory, and only - additional GPUs should have the overhead factor applied. + The first GPU keeps its full free memory; the overhead factor applies + only to additional GPUs. """ def _make_utilization(self, devices): diff --git a/studio/backend/tests/test_host_defaults.py b/studio/backend/tests/test_host_defaults.py index 8b81474e92..5c7129bc65 100644 --- a/studio/backend/tests/test_host_defaults.py +++ b/studio/backend/tests/test_host_defaults.py @@ -20,10 +20,7 @@ def _parse_function_param_defaults(source: str, func_name: str) -> dict: """ tree = ast.parse(source) for node in ast.walk(tree): - if ( - isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and node.name == func_name - ): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: result = {} all_args = node.args.args defaults = node.args.defaults @@ -38,10 +35,10 @@ def _parse_function_param_defaults(source: str, func_name: str) -> dict: def _parse_argparse_add_argument_default(source: str, option_name: str): - """Return the 'default' kwarg value for add_argument(option_name, ...) in *source*. + """Return the 'default' kwarg for add_argument(option_name, ...) in *source*. - Walks the entire module so the call can live in __main__ or in a helper - function — only handles ast.Constant defaults. + Walks the whole module so the call may live in __main__ or a helper; + only handles ast.Constant defaults. """ tree = ast.parse(source) for node in ast.walk(tree): @@ -62,18 +59,14 @@ def _parse_argparse_add_argument_default(source: str, option_name: str): def test_run_server_default_host_is_loopback(): - """run_server() parameter default for 'host' must be 127.0.0.1, not 0.0.0.0. + """run_server() 'host' default must be 127.0.0.1, not 0.0.0.0. - Binding to 0.0.0.0 by default exposes the service on all network - interfaces, contradicting the documented "privacy first / 100% local" - guarantee. Loopback (127.0.0.1) is the least-permissive default; - users who need network access can pass -H 0.0.0.0 explicitly. + 0.0.0.0 exposes the service on all interfaces; loopback is the + least-permissive default. Users needing network access pass -H 0.0.0.0. """ source = _RUN_PY.read_text() defaults = _parse_function_param_defaults(source, "run_server") - assert ( - "host" in defaults - ), "run_server() must have a 'host' parameter with a default" + assert "host" in defaults, "run_server() must have a 'host' parameter with a default" host_default = defaults["host"] assert host_default == "127.0.0.1", ( f"run_server() host default must be '127.0.0.1' (loopback) " @@ -86,13 +79,11 @@ def test_argparse_default_host_is_loopback(): """argparse --host add_argument default must be 127.0.0.1. When run.py is invoked directly (python run.py), the argparse default - should match the function default so direct execution is equally safe. + must match the function default so direct execution is equally safe. """ source = _RUN_PY.read_text() host_default = _parse_argparse_add_argument_default(source, "--host") - assert ( - host_default is not None - ), "Could not find add_argument('--host', ...) in run.py" + assert host_default is not None, "Could not find add_argument('--host', ...) in run.py" assert ( host_default == "127.0.0.1" ), f"run.py argparse --host default must be '127.0.0.1', got '{host_default}'" diff --git a/studio/backend/tests/test_index_bootstrap_origin.py b/studio/backend/tests/test_index_bootstrap_origin.py index 89f7613ee4..b9d7f58867 100644 --- a/studio/backend/tests/test_index_bootstrap_origin.py +++ b/studio/backend/tests/test_index_bootstrap_origin.py @@ -2,8 +2,9 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Regression coverage for the bootstrap-pw cross-origin leak (PR 5739). -``_is_same_origin_request`` gates ``_inject_bootstrap`` so the seeded -admin password only ships to same-origin callers. + +``_is_same_origin_request`` gates ``_inject_bootstrap`` so the seeded admin +password only ships to same-origin callers. """ import os @@ -13,7 +14,11 @@ from unittest.mock import MagicMock import pytest -def _build_request(host: str, origin: str | None, scheme: str = "http") -> MagicMock: +def _build_request( + host: str, + origin: str | None, + scheme: str = "http", +) -> MagicMock: request = MagicMock() request.url.scheme = scheme request.url.netloc = host @@ -23,21 +28,18 @@ def _build_request(host: str, origin: str | None, scheme: str = "http") -> Magic def test_is_same_origin_request_missing_origin_is_same_origin(monkeypatch): from main import _is_same_origin_request - req = _build_request("127.0.0.1:8888", origin = None) assert _is_same_origin_request(req) is True def test_is_same_origin_request_matching_origin_is_same_origin(): from main import _is_same_origin_request - req = _build_request("127.0.0.1:8888", origin = "http://127.0.0.1:8888") assert _is_same_origin_request(req) is True def test_is_same_origin_request_evil_origin_is_cross_origin(): from main import _is_same_origin_request - req = _build_request("127.0.0.1:8888", origin = "https://evil.example") assert _is_same_origin_request(req) is False @@ -45,7 +47,6 @@ def test_is_same_origin_request_evil_origin_is_cross_origin(): def test_is_same_origin_request_scheme_mismatch_is_cross_origin(): # https origin against an http listener is not same-origin. from main import _is_same_origin_request - req = _build_request("127.0.0.1:8888", origin = "https://127.0.0.1:8888") assert _is_same_origin_request(req) is False @@ -53,7 +54,6 @@ def test_is_same_origin_request_scheme_mismatch_is_cross_origin(): def test_is_same_origin_request_port_mismatch_is_cross_origin(): # Same host different port is not same-origin per the web platform. from main import _is_same_origin_request - req = _build_request("127.0.0.1:8888", origin = "http://127.0.0.1:5173") assert _is_same_origin_request(req) is False @@ -62,20 +62,15 @@ def test_is_same_origin_request_port_mismatch_is_cross_origin(): def test_is_same_origin_request_https_default_port_stripped_on_origin(): - """RFC 6454 strips default ports on Origin; Starlette's netloc may still - carry ``:443``. Canonicalise both sides so this stays same-origin. - """ + """RFC 6454 strips default ports on Origin; canonicalise both sides so this stays same-origin.""" from main import _is_same_origin_request - req = _build_request( - "example.com:443", origin = "https://example.com", scheme = "https" - ) + req = _build_request("example.com:443", origin = "https://example.com", scheme = "https") assert _is_same_origin_request(req) is True def test_is_same_origin_request_http_default_port_stripped_on_origin(): from main import _is_same_origin_request - req = _build_request("example.com:80", origin = "http://example.com") assert _is_same_origin_request(req) is True @@ -84,9 +79,7 @@ def test_is_same_origin_request_default_port_present_on_origin(): """Mirror case: Origin carries the default port, netloc doesn't. Same-origin.""" from main import _is_same_origin_request - req = _build_request( - "example.com", origin = "https://example.com:443", scheme = "https" - ) + req = _build_request("example.com", origin = "https://example.com:443", scheme = "https") assert _is_same_origin_request(req) is True @@ -115,9 +108,7 @@ def test_is_same_origin_request_null_origin_is_cross_origin(): def test_is_same_origin_request_unparseable_origin_is_cross_origin(): - """Garbage values without a host fall to cross-origin; a malformed header - must not leak the bootstrap. - """ + """Hostless garbage falls to cross-origin so a malformed header can't leak the bootstrap.""" from main import _is_same_origin_request req = _build_request("example.com", origin = "not-a-url") @@ -125,9 +116,7 @@ def test_is_same_origin_request_unparseable_origin_is_cross_origin(): def test_is_same_origin_request_userinfo_in_netloc_ignored(): - """``user:pass@host:port`` netlocs (RFC 3986) must compare equal to the - credentials-less Origin. - """ + """``user:pass@host:port`` netlocs (RFC 3986) must compare equal to the credentials-less Origin.""" from main import _is_same_origin_request req = _build_request("user:pass@example.com:80", origin = "http://example.com") @@ -138,7 +127,5 @@ def test_is_same_origin_request_explicit_non_default_port_still_mismatch(): """Canonicalisation does NOT collapse non-default ports to default.""" from main import _is_same_origin_request - req = _build_request( - "example.com", origin = "https://example.com:9999", scheme = "https" - ) + req = _build_request("example.com", origin = "https://example.com:9999", scheme = "https") assert _is_same_origin_request(req) is False diff --git a/studio/backend/tests/test_index_bootstrap_origin_extra.py b/studio/backend/tests/test_index_bootstrap_origin_extra.py index aea6b36a96..feda88c14c 100644 --- a/studio/backend/tests/test_index_bootstrap_origin_extra.py +++ b/studio/backend/tests/test_index_bootstrap_origin_extra.py @@ -2,15 +2,19 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Extra edge-case coverage for the bootstrap-pw cross-origin gate. -Companion to ``test_index_bootstrap_origin.py``: IPv6 netlocs, opaque -origins (``data:``, ``blob:``), comma-joined multi-Origin headers, and -the ``localhost`` vs ``127.0.0.1`` distinct-origin rule. +Companion to ``test_index_bootstrap_origin.py``: IPv6 netlocs, opaque origins +(``data:``, ``blob:``), comma-joined multi-Origin headers, and the +``localhost`` vs ``127.0.0.1`` distinct-origin rule. """ from unittest.mock import MagicMock -def _build_request(host: str, origin, scheme: str = "http") -> MagicMock: +def _build_request( + host: str, + origin, + scheme: str = "http", +) -> MagicMock: request = MagicMock() request.url.scheme = scheme request.url.netloc = host @@ -24,7 +28,7 @@ def _build_request(host: str, origin, scheme: str = "http") -> MagicMock: def test_is_same_origin_request_ipv6_loopback_same_origin(): """Studio supports ``-H ::1`` binds; netloc is ``[::1]:8902``. Bare ``partition(":")`` mis-parses the bracketed form and would refuse the - bootstrap on legitimate same-origin nav. + bootstrap on legitimate same-origin navigation. """ from main import _is_same_origin_request @@ -34,7 +38,6 @@ def test_is_same_origin_request_ipv6_loopback_same_origin(): def test_is_same_origin_request_ipv6_full_address_same_origin(): from main import _is_same_origin_request - req = _build_request( "[2001:db8::1]:8443", origin = "https://[2001:db8::1]:8443", @@ -65,21 +68,18 @@ def test_is_same_origin_request_ipv6_case_insensitive(): def test_is_same_origin_request_ipv6_different_host_cross_origin(): from main import _is_same_origin_request - req = _build_request("[::1]:8902", origin = "http://[2001:db8::1]:8902") assert _is_same_origin_request(req) is False def test_is_same_origin_request_ipv6_port_mismatch_cross_origin(): from main import _is_same_origin_request - req = _build_request("[::1]:8902", origin = "http://[::1]:9999") assert _is_same_origin_request(req) is False def test_is_same_origin_request_ipv6_userinfo_stripped(): from main import _is_same_origin_request - req = _build_request("user:pass@[::1]:8902", origin = "http://[::1]:8902") assert _is_same_origin_request(req) is True @@ -88,21 +88,15 @@ def test_is_same_origin_request_ipv6_userinfo_stripped(): def test_is_same_origin_request_data_url_origin_is_cross_origin(): - """``data:`` URLs are opaque origins (HTML living standard); no host, - never same-origin. - """ + """``data:`` URLs are opaque origins (HTML living standard); no host, never same-origin.""" from main import _is_same_origin_request - req = _build_request( - "127.0.0.1:8902", origin = "data:text/html," - ) + req = _build_request("127.0.0.1:8902", origin = "data:text/html,") assert _is_same_origin_request(req) is False def test_is_same_origin_request_blob_url_origin_is_cross_origin(): - """``blob:`` URLs carry the inner origin only in non-canonical form; the - canonical comparison rejects them. - """ + """``blob:`` URLs carry the inner origin only in non-canonical form; the canonical comparison rejects them.""" from main import _is_same_origin_request req = _build_request("127.0.0.1:8902", origin = "blob:http://127.0.0.1:8902/uuid") @@ -110,8 +104,8 @@ def test_is_same_origin_request_blob_url_origin_is_cross_origin(): def test_is_same_origin_request_file_url_origin_is_cross_origin(): - """``file://`` pages usually send ``Origin: null``; historical engines - sent ``Origin: file://``. Neither is same-origin vs an http listener. + """``file://`` pages usually send ``Origin: null``; older engines sent + ``Origin: file://``. Neither is same-origin vs an http listener. """ from main import _is_same_origin_request @@ -123,8 +117,8 @@ def test_is_same_origin_request_file_url_origin_is_cross_origin(): def test_is_same_origin_request_comma_joined_origins_cross_origin(): - """Starlette concatenates repeated headers with ``, ``; the canonical - parser can't safely split this, so it falls to cross-origin. + """Starlette joins repeated headers with ``, ``; the canonical parser can't + safely split this, so it falls to cross-origin. """ from main import _is_same_origin_request @@ -139,8 +133,8 @@ def test_is_same_origin_request_comma_joined_origins_cross_origin(): def test_is_same_origin_request_localhost_vs_127_is_cross_origin(): - """Browsers treat ``localhost`` and ``127.0.0.1`` as distinct origins; - the canonical comparison must not DNS-collapse them. + """Browsers treat ``localhost`` and ``127.0.0.1`` as distinct origins; the + canonical comparison must not DNS-collapse them. """ from main import _is_same_origin_request @@ -150,7 +144,6 @@ def test_is_same_origin_request_localhost_vs_127_is_cross_origin(): def test_is_same_origin_request_127_vs_localhost_is_cross_origin(): from main import _is_same_origin_request - req = _build_request("localhost:8902", origin = "http://127.0.0.1:8902") assert _is_same_origin_request(req) is False @@ -160,7 +153,7 @@ def test_is_same_origin_request_127_vs_localhost_is_cross_origin(): def test_is_same_origin_request_malformed_ipv6_bracket_is_cross_origin(): """``urlparse`` raises ``ValueError('Invalid IPv6 URL')`` on unclosed - brackets (CVE-2024-11168 hardening). The gate must swallow and fall to + brackets (CVE-2024-11168 hardening). The gate must swallow it and fall to cross-origin rather than 500 the SPA handler. """ from main import _is_same_origin_request @@ -187,8 +180,8 @@ def test_is_same_origin_request_bracket_with_trailing_garbage_is_cross_origin(): def test_is_same_origin_request_empty_origin_header_is_cross_origin(): - """Explicit empty ``Origin:`` is not a valid serialised origin and must - not be conflated with a missing header; cross-origin, bootstrap withheld. + """Explicit empty ``Origin:`` is not a valid serialised origin and must not + be conflated with a missing header; cross-origin, bootstrap withheld. """ from main import _is_same_origin_request diff --git a/studio/backend/tests/test_inference_model_validation.py b/studio/backend/tests/test_inference_model_validation.py index ebd9c6c722..2427ad35fa 100644 --- a/studio/backend/tests/test_inference_model_validation.py +++ b/studio/backend/tests/test_inference_model_validation.py @@ -170,16 +170,15 @@ def test_walkback_does_not_cross_user_turn(): ] ) last = req.messages[-1].tool_call_id - # The walkback must NOT pick old_call because a user turn intervenes; - # falls back to synth. + # Walkback must NOT pick old_call across a user turn; falls back to synth. assert last is not None assert last != "old_call" assert last.startswith("call_") def test_walkback_skips_explicitly_consumed_tool_call_id(): - """Sibling tool result with an explicit id must reserve its assistant - slot so a follow-up missing-id result picks the OTHER tool call.""" + """An explicit-id tool result reserves its assistant slot so a + follow-up missing-id result picks the OTHER tool call.""" req = _req( [ { @@ -202,15 +201,12 @@ def test_walkback_skips_explicitly_consumed_tool_call_id(): {"role": "tool", "content": "second result"}, ] ) - assert [m.tool_call_id for m in req.messages if m.role == "tool"] == [ - "call_a", - "call_b", - ] + assert [m.tool_call_id for m in req.messages if m.role == "tool"] == ["call_a", "call_b"] def test_walkback_handles_malformed_function_string(): """A tool_call with ``function`` as a string (provider quirk) must not - raise; resolution falls back to fallback id selection.""" + raise; resolution falls back to id selection.""" req = _req( [ { diff --git a/studio/backend/tests/test_inference_orchestrator_crash_message.py b/studio/backend/tests/test_inference_orchestrator_crash_message.py new file mode 100644 index 0000000000..be1a673e62 --- /dev/null +++ b/studio/backend/tests/test_inference_orchestrator_crash_message.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from pathlib import Path +from types import SimpleNamespace +import importlib.util +import sys + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + + +def test_subprocess_crash_message_includes_signal_and_oom_hint(): + spec = importlib.util.spec_from_file_location( + "inference_orchestrator_under_test", + Path(__file__).resolve().parent.parent / "core/inference/orchestrator.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + orchestrator = module.InferenceOrchestrator.__new__(module.InferenceOrchestrator) + orchestrator._proc = SimpleNamespace(pid = 1234, exitcode = -9) + + msg = orchestrator._subprocess_crash_message("wait") + + assert msg.startswith("The inference worker stopped unexpectedly while loading the model.") + assert "memory pressure" in msg + assert "smaller model" in msg + assert "Details:" in msg + assert "signal=SIGKILL" in msg + assert "exitcode=-9" in msg diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py new file mode 100644 index 0000000000..ede5629664 --- /dev/null +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode. + +These back the in-app update for source-build (markerless) installs: the backend +asks the installer whether an official prebuilt exists for this host without +downloading. Network and host detection are stubbed; no GPU or internet needed. +""" + +from __future__ import annotations + +import importlib +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_studio = Path(__file__).resolve().parent.parent.parent +if str(_studio) not in sys.path: + sys.path.insert(0, str(_studio)) + +ilp = importlib.import_module("install_llama_prebuilt") + +if not hasattr(ilp, "published_repo_for_host") or not hasattr( + ilp, "resolve_simple_install_release_plans" +): + pytest.skip("PR symbols not present - check branch", allow_module_level = True) + +FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp +UPSTREAM = ilp.UPSTREAM_REPO # ggml-org/llama.cpp + + +def _host(**kw): + base = dict( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = False, + is_macos = False, + is_x86_64 = False, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = False, + rocm_gfx_target = None, + macos_version = None, + ) + base.update(kw) + return ilp.HostInfo(**base) + + +def test_published_repo_for_host(): + # CPU-only Linux (x64 and arm64) -> ggml-org upstream. + assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM + assert ( + ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) + == UPSTREAM + ) + # GPU Linux -> fork. + assert ( + ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) + == FORK + ) + assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK + # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). + assert ( + ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) + == UPSTREAM + ) + # GPU Windows -> fork. + assert ( + ilp.published_repo_for_host( + _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) + ) + == FORK + ) + # macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS). + assert ( + ilp.published_repo_for_host( + _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") + ) + == FORK + ) + # Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling). + assert ( + ilp.published_repo_for_host( + _host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True + ) + == FORK + ) + # The tooling hint is Linux-only: Windows CPU stays on ggml-org. + assert ( + ilp.published_repo_for_host( + _host(system = "Windows", is_windows = True, is_x86_64 = True), + linux_amd_tooling_present = True, + ) + == UPSTREAM + ) + + +def _run_resolve(monkeypatch, capsys, plans_or_exc): + monkeypatch.setattr( + ilp, + "detect_host", + lambda: _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64"), + ) + + def _resolver(tag, host, repo, published_release_tag): + if isinstance(plans_or_exc, Exception): + raise plans_or_exc + return ("b9585", plans_or_exc) + + monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"], + ) + rc = ilp.main() + assert rc == ilp.EXIT_SUCCESS + return json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + + +def test_resolve_prebuilt_available(monkeypatch, capsys): + plan = SimpleNamespace( + release_tag = "b9585", + llama_tag = "b9585", + attempts = [ + SimpleNamespace(name = "llama-b9585-bin-macos-arm64.tar.gz", install_kind = "macos-arm64") + ], + ) + out = _run_resolve(monkeypatch, capsys, [plan]) + assert out["prebuilt_available"] is True + assert out["repo"] == FORK + assert out["release_tag"] == "b9585" + assert out["asset"] == "llama-b9585-bin-macos-arm64.tar.gz" + assert out["install_kind"] == "macos-arm64" + + +def test_resolve_prebuilt_unavailable(monkeypatch, capsys): + out = _run_resolve(monkeypatch, capsys, ilp.PrebuiltFallback("no macOS asset")) + assert out["prebuilt_available"] is False + assert out["repo"] == FORK + + +def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): + # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the + # fork so a HIP source build is not offered an upstream CPU prebuilt. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") + seen = {} + + def _resolver(tag, host, repo, published_release_tag): + seen["repo"] = repo + raise ilp.PrebuiltFallback("no asset") + + monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"], + ) + assert ilp.main() == ilp.EXIT_SUCCESS + out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + assert seen["repo"] == FORK + assert out["repo"] == FORK diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index d52a58a25c..cd834b345b 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -7,8 +7,7 @@ Covers the GGUF metadata parser, _can_estimate_kv gate, all 5 estimation paths (MLA, Hybrid Mamba, Sliding Window, Standard GQA, Legacy), KV cache quantization, edge cases, and lifecycle (init/unload/reparse). -Requires no GPU, network, or external libraries beyond pytest. -Cross-platform: Linux, macOS, Windows, WSL. +No GPU, network, or libraries beyond pytest. Cross-platform. """ import io @@ -20,10 +19,8 @@ from pathlib import Path import pytest -# --------------------------------------------------------------------------- -# Stub heavy / unavailable external dependencies before importing the -# module under test. Same pattern as test_native_context_length.py. -# --------------------------------------------------------------------------- +# Stub heavy / unavailable deps before importing the module under test. +# Same pattern as test_native_context_length.py. _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: @@ -38,11 +35,9 @@ sys.modules.setdefault("loggers", _loggers_stub) _structlog_stub = _types.ModuleType("structlog") sys.modules.setdefault("structlog", _structlog_stub) -# httpx -- only stub when the real library isn't installed. Stubbing -# unconditionally would shadow ``HTTPError`` / ``Response`` etc. that -# ``huggingface_hub.errors`` imports at module load time, which causes -# the transformers introspection tier to silently return None inside -# the test process. +# httpx -- only stub when the real library is missing. Unconditional stubbing +# shadows HTTPError/Response that huggingface_hub.errors imports at load time, +# silently breaking the transformers introspection tier. try: import httpx as _httpx_real # noqa: F401 except ImportError: @@ -78,15 +73,13 @@ except ImportError: from core.inference.llama_cpp import LlamaCppBackend -# --------------------------------------------------------------------------- # Helpers -# --------------------------------------------------------------------------- def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: - """Build a minimal GGUF v3 binary blob with the given KV metadata. + """Build a minimal GGUF v3 blob with the given KV metadata. - Supports the scalar and simple array metadata used by the parser. + Supports the scalar and simple array metadata the parser uses. """ buf = io.BytesIO() # Header: magic, version, tensor_count, kv_count @@ -128,13 +121,14 @@ def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: def _backend_from_gguf( - arch: str, fields: dict, general: dict | None = None + arch: str, + fields: dict, + general: dict | None = None, ) -> LlamaCppBackend: """Create a LlamaCppBackend with parsed GGUF metadata from given fields. - `general` lets a test inject extra `general.*` metadata (used to - verify the dynamic SWA resolver picks up source-repo hints from - GGUFs that ship them). + `general` injects extra `general.*` metadata, to verify the dynamic + SWA resolver picks up source-repo hints from GGUFs that ship them. """ kv = {"general.architecture": arch} for k, v in (general or {}).items(): @@ -155,13 +149,11 @@ def _backend_from_gguf( os.unlink(path) -# --------------------------------------------------------------------------- # A. GGUF Parser Tests -# --------------------------------------------------------------------------- class TestGGUFParserNewFields: - """Verify that architecture-aware fields are correctly parsed.""" + """Architecture-aware fields are parsed correctly.""" @pytest.mark.parametrize( "field,gguf_key,value", @@ -215,9 +207,9 @@ class TestGGUFParserNewFields: ) # Per-layer KV head count is preserved exactly... assert b._n_kv_heads_by_layer == [8, 8, 8, 8, 8, 2] - # ...and mirrored into the scalar field as a conservative max so - # non-SWA estimator paths and any caller using - # `n_kv = self._n_kv_heads or ...` get a safe upper bound. + # ...and mirrored into the scalar field as a conservative max, so + # non-SWA paths and callers using `n_kv = self._n_kv_heads or ...` + # get a safe upper bound. assert b._n_kv_heads == 8 assert b._sliding_window_pattern == [True, True, True, True, True, False] @@ -268,7 +260,7 @@ class TestArchSwaPatternDefaults: assert b._sliding_window_pattern is None def test_explicit_pattern_overrides_arch_default(self): - # Period=6 is the gemma3 default; the explicit array must win. + # gemma3 default is period=6; the explicit array must win. b = _backend_from_gguf( "gemma3", { @@ -308,8 +300,8 @@ class TestArchSwaPatternDefaults: "arch", ["llama", "qwen2", "qwen3", "mistral", "mistral3", "glm4", "llama4"] ) def test_non_swa_arch_uses_full_attention_path(self, arch): - # Pure-GQA arches: GGUF has no sliding_window, no synthetic - # pattern, estimator hits Path 4. + # Pure-GQA arches: no sliding_window, no synthetic pattern, + # estimator hits Path 4. b = _backend_from_gguf( arch, { @@ -338,7 +330,7 @@ class TestArchSwaPatternDefaults: "embedding_length": 5376, } with_default = _backend_from_gguf("gemma3", common) - # Arch not in the table -> legacy 1/4 path. + # Arch not in table -> legacy 1/4 path. without_default = _backend_from_gguf("totallymadeupv7", common) kv_default = with_default._estimate_kv_cache_bytes(131072, "f16") @@ -346,8 +338,7 @@ class TestArchSwaPatternDefaults: assert kv_default > 0 assert kv_legacy > 0 assert kv_default < kv_legacy, ( - f"arch fallback should under-shoot legacy estimate: " - f"{kv_default} >= {kv_legacy}" + f"arch fallback should under-shoot legacy estimate: " f"{kv_default} >= {kv_legacy}" ) def test_scalar_sliding_window_pattern_expanded(self): @@ -429,26 +420,13 @@ class TestDynamicSwaResolver: def test_period_from_layer_types_finds_smallest_period(self): from core.inference.llama_cpp import _period_from_layer_types - # gemma3 (1 global per 6), gpt-oss (alternating), gemma3n (1 per 5). - assert ( - _period_from_layer_types( - (["sliding_attention"] * 5 + ["full_attention"]) * 4 - ) - == 6 - ) - assert ( - _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2 - ) - assert ( - _period_from_layer_types( - (["sliding_attention"] * 4 + ["full_attention"]) * 7 - ) - == 5 - ) + # gemma3 (1 global/6), gpt-oss (alternating), gemma3n (1/5). + assert _period_from_layer_types((["sliding_attention"] * 5 + ["full_attention"]) * 4) == 6 + assert _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2 + assert _period_from_layer_types((["sliding_attention"] * 4 + ["full_attention"]) * 7) == 5 def test_period_from_layer_types_returns_none_for_aperiodic(self): from core.inference.llama_cpp import _period_from_layer_types - lt = [ "sliding_attention", "full_attention", @@ -469,9 +447,7 @@ class TestDynamicSwaResolver: == "google/gemma-3-1b-it" ) assert ( - _hf_repo_from_url( - "https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json" - ) + _hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json") == "google/gemma-3-1b-it" ) for bad in [ @@ -495,14 +471,14 @@ class TestDynamicSwaResolver: def test_disk_cache_takes_precedence_over_bootstrap(self, monkeypatch, tmp_path): self._isolate_cache(monkeypatch, tmp_path) - # Override bootstrap=6 with a cached period=3. + # Cached period=3 overrides bootstrap=6. with open(tmp_path / "swa_cache.json", "w") as f: json.dump({"gemma3": 3}, f) b = _backend_from_gguf("gemma3", dict(_SWA_FIELDS, block_count = 18)) assert b._sliding_window_pattern == [(i + 1) % 3 != 0 for i in range(18)] def test_disk_cache_supports_array_entries(self, monkeypatch, tmp_path): - # Aperiodic mask gets tiled across n_layers. + # Aperiodic mask is tiled across n_layers. self._isolate_cache(monkeypatch, tmp_path) mask = [True, False, True, True, False, True, False, False] with open(tmp_path / "swa_cache.json", "w") as f: @@ -524,9 +500,7 @@ class TestDynamicSwaResolver: b = _backend_from_gguf( "newmodel", _SWA_FIELDS, - general = { - "general.source.huggingface.repository": "vendor/newmodel-1b-instruct" - }, + general = {"general.source.huggingface.repository": "vendor/newmodel-1b-instruct"}, ) assert b._sliding_window_pattern == [(i + 1) % 4 != 0 for i in range(12)] assert calls == ["vendor/newmodel-1b-instruct"] @@ -572,10 +546,8 @@ class TestDynamicSwaResolver: from core.inference import llama_cpp as lc monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", lambda repo_id: None) - # Force the failure into the Tier 3 path; bypass Tier 2.5. - monkeypatch.setattr( - lc, "_resolve_swa_entry_from_transformers", lambda arch: None - ) + # Force failure into Tier 3; bypass Tier 2.5. + monkeypatch.setattr(lc, "_resolve_swa_entry_from_transformers", lambda arch: None) b = _backend_from_gguf( "newmodel", _SWA_FIELDS, @@ -621,18 +593,14 @@ class TestTransformersIntrospection: class _FakeLazyMapping(dict): def __getitem__(self, k): - return ( - _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k) - ) + return _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k) import sys, types as _types fake_auto = _types.ModuleType("transformers.models.auto.configuration_auto") fake_auto.CONFIG_MAPPING_NAMES = {"brokenarch": "FakeBroken"} fake_auto.CONFIG_MAPPING = _FakeLazyMapping({"brokenarch": "FakeBroken"}) - monkeypatch.setitem( - sys.modules, "transformers.models.auto.configuration_auto", fake_auto - ) + monkeypatch.setitem(sys.modules, "transformers.models.auto.configuration_auto", fake_auto) assert lc._resolve_swa_entry_from_transformers("brokenarch") == 7 def test_returns_none_when_transformers_unavailable(self, monkeypatch): @@ -658,13 +626,10 @@ class TestTransformersIntrospection: def test_returns_none_for_arch_unknown_to_transformers(self): from core.inference.llama_cpp import _resolve_swa_entry_from_transformers - assert _resolve_swa_entry_from_transformers("totally-fake-arch-xyz") is None - def test_full_resolver_uses_transformers_before_hf_fetch( - self, monkeypatch, tmp_path - ): - # With bootstrap empty, Tier 2.5 must answer before Tier 3 fires. + def test_full_resolver_uses_transformers_before_hf_fetch(self, monkeypatch, tmp_path): + # Bootstrap empty: Tier 2.5 must answer before Tier 3 fires. self._isolate_cache(monkeypatch, tmp_path) from core.inference import llama_cpp as lc @@ -685,10 +650,10 @@ class TestTransformersIntrospection: class TestGGUFParserReset: - """Verify that fields are properly reset between parses.""" + """Fields are reset between parses.""" def test_reset_between_parses(self): - # First parse with all fields + # First parse: all fields set b = _backend_from_gguf( "arch1", { @@ -710,7 +675,7 @@ class TestGGUFParserReset: assert b._kv_value_length_swa == 64 assert b._ssm_inner_size == 4096 - # Second parse without those fields -- they should be None + # Second parse without those fields -- they must be None kv = {"general.architecture": "arch2", "arch2.block_count": 64} import tempfile, os @@ -732,13 +697,11 @@ class TestGGUFParserReset: assert b._n_layers == 64 -# --------------------------------------------------------------------------- # B. _can_estimate_kv Gate Tests -# --------------------------------------------------------------------------- class TestCanEstimateKV: - """Verify gate logic for all field combinations.""" + """Gate logic for all field combinations.""" def test_no_layers_returns_false(self): b = LlamaCppBackend() @@ -754,7 +717,7 @@ class TestCanEstimateKV: assert b._can_estimate_kv() def test_key_length_alone_insufficient(self): - """key_length without value_length should NOT be enough.""" + """key_length without value_length is NOT enough.""" b = LlamaCppBackend() b._n_layers = 32 b._kv_key_length = 128 @@ -792,9 +755,7 @@ class TestCanEstimateKV: assert not b._can_estimate_kv() -# --------------------------------------------------------------------------- # C. Path 1: MLA Estimation -# --------------------------------------------------------------------------- class TestMLAEstimation: @@ -824,33 +785,33 @@ class TestMLAEstimation: assert b._estimate_kv_cache_bytes(163840, "f16") == expected def test_mla_ignores_value_length(self): - """MLA should NOT add value_length -- V is reconstructed from the latent.""" + """MLA must NOT add value_length -- V is reconstructed from the latent.""" b = self._mla_backend() result = b._estimate_kv_cache_bytes(1000, "f16") - # Should be n_layers * ctx * 1 * key_len(576) * 2 + # n_layers * ctx * 1 * key_len(576) * 2 expected = 61 * 1000 * 1 * 576 * 2 assert result == expected def test_mla_fallback_when_no_key_length(self): - """If key_length is missing, fallback to kv_lora_rank + key_length_mla.""" + """No key_length: fall back to kv_lora_rank + key_length_mla.""" b = self._mla_backend(_kv_key_length = None) - # _key_length_mla=192 in default, so rope_dim=192 + # default _key_length_mla=192, so rope_dim=192 result = b._estimate_kv_cache_bytes(1000, "f16") expected = 61 * 1000 * 1 * (512 + 192) * 2 # 704 assert result == expected def test_mla_fallback_no_key_length_mla(self): - """If both key_length and key_length_mla are missing, fallback to +64.""" + """No key_length and no key_length_mla: fall back to +64.""" b = self._mla_backend(_kv_key_length = None, _key_length_mla = None) result = b._estimate_kv_cache_bytes(1000, "f16") expected = 61 * 1000 * 1 * (512 + 64) * 2 # 576 assert result == expected def test_mla_defaults_n_kv_to_1_when_heads_absent(self): - """MLA should use n_kv=1 even if n_kv_heads is None (not n_heads).""" + """MLA uses n_kv=1 even if n_kv_heads is None (not n_heads).""" b = self._mla_backend(_n_kv_heads = None) # n_heads=128 still set result = b._estimate_kv_cache_bytes(1000, "f16") - # Should use n_kv_mla=1, NOT n_heads=128 + # Uses n_kv_mla=1, NOT n_heads=128 expected = 61 * 1000 * 1 * 576 * 2 assert result == expected @@ -863,9 +824,7 @@ class TestMLAEstimation: assert result_q4 == int(61 * 1000 * 1 * 576 * 0.5625) -# --------------------------------------------------------------------------- # D. Path 2: Hybrid Mamba Estimation -# --------------------------------------------------------------------------- class TestHybridMambaEstimation: @@ -908,14 +867,14 @@ class TestHybridMambaEstimation: assert b._estimate_kv_cache_bytes(262144, "f16") == expected def test_hybrid_without_explicit_dims(self): - """Fallback to head_dim when key_length/value_length are missing.""" + """Fall back to head_dim when key_length/value_length are missing.""" b = self._hybrid_backend(_kv_key_length = None, _kv_value_length = None) head_dim = 5120 // 24 # 213 expected = 16 * 4096 * 4 * 2 * head_dim * 2 assert b._estimate_kv_cache_bytes(4096, "f16") == expected def test_fai_zero_safety(self): - """full_attention_interval=0 should not cause ZeroDivisionError.""" + """full_attention_interval=0 must not ZeroDivisionError.""" b = self._hybrid_backend(_full_attention_interval = 0) result = b._estimate_kv_cache_bytes(4096, "f16") # fai=0 -> n_attn = n_layers (all layers) @@ -923,9 +882,7 @@ class TestHybridMambaEstimation: assert result == expected -# --------------------------------------------------------------------------- # E. Path 3: Sliding Window Estimation -# --------------------------------------------------------------------------- class TestSlidingWindowEstimation: @@ -1003,7 +960,7 @@ class TestSlidingWindowEstimation: assert b._estimate_kv_cache_bytes(ctx, "f16") == expected(ctx) def test_ctx_smaller_than_window(self): - """When context < 2 * sliding_window, SWA cache caps at ctx.""" + """When ctx < 2 * sliding_window, SWA cache caps at ctx.""" b = self._swa_backend(_sliding_window = 8192) n_global = max(1, 62 // 4) # 15 n_swa = 62 - n_global # 47 @@ -1021,9 +978,7 @@ class TestSlidingWindowEstimation: assert b._estimate_kv_cache_bytes(1000, "f16") == expected -# --------------------------------------------------------------------------- # F. Path 4: Standard GQA Estimation -# --------------------------------------------------------------------------- class TestStandardGQAEstimation: @@ -1056,20 +1011,18 @@ class TestStandardGQAEstimation: assert b._estimate_kv_cache_bytes(4096, "f16") == expected def test_differs_from_legacy(self): - """GQA path should differ from legacy when key_length != embed//n_heads.""" + """GQA path differs from legacy when key_length != embed//n_heads.""" b = self._gqa_backend() head_dim = 1024 // 16 # 64 gqa_result = b._estimate_kv_cache_bytes(4096, "f16") - # Legacy would use: 2 * 8 * 64 * 28 * 4096 * 2 + # Legacy: 2 * 8 * 64 * 28 * 4096 * 2 legacy_result = int(2 * 8 * head_dim * 28 * 4096 * 2) # GQA: 28 * 4096 * 8 * (128+128) * 2 -- uses actual key_length=128 assert gqa_result != legacy_result assert gqa_result > legacy_result # key_length (128) > head_dim (64) -# --------------------------------------------------------------------------- # G. Path 5: Legacy Fallback Estimation -# --------------------------------------------------------------------------- class TestLegacyEstimation: @@ -1102,7 +1055,7 @@ class TestLegacyEstimation: assert b._estimate_kv_cache_bytes(4096, "f16") == expected def test_legacy_identical_to_old_formula(self): - """Confirm legacy path produces the same result as the pre-PR formula.""" + """Legacy path matches the pre-PR formula.""" b = self._legacy_backend() n_layers = 32 n_kv_heads = 8 @@ -1113,16 +1066,14 @@ class TestLegacyEstimation: assert b._estimate_kv_cache_bytes(n_ctx, "f16") == old_formula -# --------------------------------------------------------------------------- # H. Path Priority (selection order) -# --------------------------------------------------------------------------- class TestPathPriority: """Confirm: MLA > Hybrid Mamba > SWA > GQA > Legacy.""" def test_mla_takes_priority_over_all(self): - """If kv_lora_rank is set, MLA path is used even if other fields are present.""" + """If kv_lora_rank is set, MLA path wins even with other fields present.""" b = LlamaCppBackend() b._n_layers = 61 b._n_kv_heads = 1 @@ -1157,9 +1108,9 @@ class TestPathPriority: assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid def test_all_paths_produce_different_values(self): - """With carefully chosen params, each path should yield a distinct value.""" - # Use embedding_length=768 so legacy head_dim (768//16=48) differs from - # key_length (256), and MLA key_len (256) != legacy K+V (2*48=96). + """With chosen params, each path yields a distinct value.""" + # embedding_length=768 so legacy head_dim (768//16=48) != key_length + # (256), and MLA key_len (256) != legacy K+V (2*48=96). params = { "_n_layers": 40, "_n_kv_heads": 4, @@ -1210,13 +1161,11 @@ class TestPathPriority: assert len(set(values)) == 5, f"Expected 5 distinct values, got {values}" -# --------------------------------------------------------------------------- # I. KV Cache Quantization -# --------------------------------------------------------------------------- class TestQuantization: - """Verify all supported cache_type_kv values produce correct scaling.""" + """All supported cache_type_kv values scale correctly.""" @pytest.mark.parametrize( "cache_type,expected_bpe", @@ -1247,9 +1196,7 @@ class TestQuantization: assert result == expected -# --------------------------------------------------------------------------- # J. Edge Cases -# --------------------------------------------------------------------------- class TestEdgeCases: @@ -1310,10 +1257,8 @@ class TestEdgeCases: assert result == expected -# --------------------------------------------------------------------------- # J2. Server-flag knobs (--swa-full, --kv-unified/--parallel, # --ctx-checkpoints, --kv-offload) -# --------------------------------------------------------------------------- class TestServerFlags: @@ -1328,8 +1273,7 @@ class TestServerFlags: "_kv_key_length": 256, "_kv_value_length": 256, "_sliding_window": 512, - "_sliding_window_pattern": [True, True, True, True, True, False] * 4 - + [True, True], + "_sliding_window_pattern": [True, True, True, True, True, False] * 4 + [True, True], } defaults.update(overrides) b = LlamaCppBackend() @@ -1358,7 +1302,7 @@ class TestServerFlags: b = self._swa_backend() ctx = 32_768 flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) - # With swa_full, every layer caches n_ctx -- equals path 4 sizing. + # swa_full: every layer caches n_ctx -- equals path 4 sizing. kv_per_token = 4 * (256 + 256) * 2 # n_kv_heads * (k+v) * f16 expected = 26 * ctx * kv_per_token assert flagged == expected @@ -1385,19 +1329,16 @@ class TestServerFlags: def test_swa_full_suppresses_checkpoint_term(self): b = self._swa_backend() with_cp = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8) - with_cp_full = b._estimate_kv_cache_bytes( - 8192, "f16", ctx_checkpoints = 8, swa_full = True - ) + with_cp_full = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8, swa_full = True) no_cp_full = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True) # Checkpoints only matter when SWA layers don't already keep n_ctx. assert with_cp_full == no_cp_full assert with_cp > b._estimate_kv_cache_bytes(8192, "f16") # ── --parallel + --kv-unified ────────────────────────────────── - # Empirically verified against llama-server: non-SWA caches partition - # n_ctx across slots (total memory constant); SWA layers are the only - # portion that scales with --parallel. --kv-unified is currently a - # no-op for memory math (kept for API forward-compat). + # Verified against llama-server: non-SWA caches partition n_ctx across + # slots (total memory constant); only SWA layers scale with --parallel. + # --kv-unified is a no-op for memory math (kept for API forward-compat). def test_gqa_kv_constant_across_parallel(self): b = self._gqa_backend() @@ -1405,9 +1346,7 @@ class TestServerFlags: for slots in (1, 2, 4, 8): for unified in (True, False): assert ( - b._estimate_kv_cache_bytes( - 4096, "f16", n_parallel = slots, kv_unified = unified - ) + b._estimate_kv_cache_bytes(4096, "f16", n_parallel = slots, kv_unified = unified) == baseline ) @@ -1416,9 +1355,7 @@ class TestServerFlags: baseline = b._estimate_kv_cache_bytes(4096, "f16") for unified in (True, False): assert ( - b._estimate_kv_cache_bytes( - 4096, "f16", n_parallel = 0, kv_unified = unified - ) + b._estimate_kv_cache_bytes(4096, "f16", n_parallel = 0, kv_unified = unified) == baseline ) @@ -1426,15 +1363,13 @@ class TestServerFlags: b = self._swa_backend() ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16") - # Decompose baseline by walking the same loop the estimator does. + # Decompose baseline by walking the estimator's own loop. swa = b._sliding_window per_token_global = 4 * (256 + 256) * 2 # n_kv * (k+v) * f16 per_token_swa = 4 * (256 + 256) * 2 # k_swa/val_swa fall back per_slot_swa_cells = min(ctx, 2 * swa) # not clamped at parallel=1 global_bytes = sum( - ctx * per_token_global - for f in b._sliding_window_pattern[: b._n_layers] - if not f + ctx * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f ) swa_bytes_per_slot = sum( per_slot_swa_cells * per_token_swa @@ -1443,18 +1378,14 @@ class TestServerFlags: ) # Sanity: parallel=1 reproduces baseline exactly assert global_bytes + swa_bytes_per_slot == baseline - # Only SWA portion scales by parallel + # Only the SWA portion scales by parallel for slots in (1, 2, 3, 4): - scaled = b._estimate_kv_cache_bytes( - ctx, "f16", n_parallel = slots, kv_unified = False - ) - # SWA cells get clamped to per_slot_ctx when ctx/slots < 2*swa + scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False) + # SWA cells clamp to per_slot_ctx when ctx/slots < 2*swa per_slot_ctx = max(1, ctx // slots) cells = min(ctx, 2 * swa, per_slot_ctx) swa_bps = sum( - cells * per_token_swa - for f in b._sliding_window_pattern[: b._n_layers] - if f + cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f ) assert scaled == global_bytes + slots * swa_bps @@ -1469,9 +1400,7 @@ class TestServerFlags: for slots in (1, 2, 4, 8): for unified in (True, False): assert ( - b._estimate_kv_cache_bytes( - 8192, "f16", n_parallel = slots, kv_unified = unified - ) + b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified) == baseline ) @@ -1492,10 +1421,8 @@ class TestServerFlags: ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16") flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4) - # 22 SWA layers * 4 checkpoints * 512 cells * 4 heads * (256+256) * 2 bytes - n_swa_layers = sum( - 1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f - ) + # 22 SWA layers * 4 cps * 512 cells * 4 heads * (256+256) * 2 bytes + n_swa_layers = sum(1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f) per_layer = 4 * 512 * 4 * (256 + 256) * 2 assert flagged == baseline + n_swa_layers * per_layer @@ -1512,7 +1439,7 @@ class TestServerFlags: def test_ctx_checkpoints_compose_with_n_parallel(self): # Only the SWA + checkpoint portion scales by n_parallel; the - # global-layer portion stays constant. + # global-layer portion is constant. b = self._swa_backend() ctx = 8192 swa = b._sliding_window @@ -1529,15 +1456,13 @@ class TestServerFlags: flagged = b._estimate_kv_cache_bytes( ctx, "f16", ctx_checkpoints = 4, n_parallel = slots, kv_unified = False ) - assert flagged == global_bytes + slots * ( - swa_bytes_per_slot + cp_extra_per_slot - ) + assert flagged == global_bytes + slots * (swa_bytes_per_slot + cp_extra_per_slot) # ── --kv-offload (kv_on_gpu) ─────────────────────────────────── def test_fit_returns_requested_when_kv_off_gpu(self): b = self._gqa_backend() - # Tiny VRAM budget -- normally would force a reduction. + # Tiny VRAM budget -- would normally force a reduction. fitted = b._fit_context_to_vram( requested_ctx = 32_768, available_mib = 1, @@ -1559,8 +1484,8 @@ class TestServerFlags: assert fitted < 32_768 def test_fit_mtp_engaged_returns_smaller_or_equal_context(self): - # MTP-engaged budget is 0.85 of available; non-MTP is 0.90. - # On a tight budget the MTP path must yield <= the non-MTP path. + # MTP budget is 0.85 of available, non-MTP is 0.90; on a tight + # budget MTP must yield <= non-MTP. b = self._gqa_backend() common = dict( requested_ctx = 32_768, @@ -1573,7 +1498,7 @@ class TestServerFlags: assert mtp <= baseline def test_fit_mtp_engaged_unchanged_when_kv_off_gpu(self): - # kv_on_gpu=False short-circuits the fit; mtp_engaged is irrelevant. + # kv_on_gpu=False short-circuits the fit; mtp_engaged irrelevant. b = self._gqa_backend() fitted = b._fit_context_to_vram( requested_ctx = 32_768, @@ -1592,7 +1517,7 @@ class TestServerFlags: kv_default = b._estimate_kv_cache_bytes(ctx, "f16") kv_full = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) assert kv_full > kv_default - # Budget = model + kv_default (rounded up) -- swa_full should not fit. + # Budget = model + kv_default (rounded up) -- swa_full must not fit. budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / 0.90 + 1 fitted_default = b._fit_context_to_vram( requested_ctx = ctx, @@ -1611,22 +1536,20 @@ class TestServerFlags: assert fitted_full < ctx -# --------------------------------------------------------------------------- # J2.5. --parallel N memory accounting (per-layer-type scaling rule) -# --------------------------------------------------------------------------- class TestParallelSWAScaling: - """Verifies the per-layer-type scaling rule against the closed form - measured from llama-server. Empirical formula on Gemma-3 270m at - ctx=8192: total_kv = 24 + parallel * 15 (MiB). + """Per-layer-type scaling rule vs the closed form measured from + llama-server. Empirical formula on Gemma-3 270m at ctx=8192: + total_kv = 24 + parallel * 15 (MiB). Rule (verified vs ``llama-server`` log on real GGUFs): * non-SWA layers: total cells = n_ctx, partitioned across slots, memory CONSTANT in n_parallel. * SWA layers: per-slot cells = 2 * sliding_window (clamped at n_ctx and at per_slot_ctx); memory LINEAR in n_parallel. - * --kv-unified is a no-op for memory math; both modes yield the + * --kv-unified is a no-op for memory math; both modes give the same total in measured cases. """ @@ -1655,9 +1578,7 @@ class TestParallelSWAScaling: "_kv_value_length": 256, "_sliding_window": 512, # 15 SWA + 3 global, mirrors gemma-3-270m - "_sliding_window_pattern": [ - t == "swa" for t in (["swa"] * 5 + ["global"]) * 3 - ], + "_sliding_window_pattern": [t == "swa" for t in (["swa"] * 5 + ["global"]) * 3], } defaults.update(overrides) b = LlamaCppBackend() @@ -1673,9 +1594,7 @@ class TestParallelSWAScaling: for slots in (1, 2, 4, 8): for unified in (True, False): assert ( - b._estimate_kv_cache_bytes( - 8192, "f16", n_parallel = slots, kv_unified = unified - ) + b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified) == baseline ) @@ -1729,9 +1648,7 @@ class TestParallelSWAScaling: cells = min(ctx, 2 * swa, per_slot_ctx) swa_bps = n_swa * cells * per_token for unified in (True, False): - got = b._estimate_kv_cache_bytes( - ctx, "f16", n_parallel = slots, kv_unified = unified - ) + got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified) assert got == global_bytes + slots * swa_bps def test_swa_fallback_scales_only_swa_portion(self): @@ -1753,7 +1670,7 @@ class TestParallelSWAScaling: def test_swa_per_slot_clamped_when_ctx_lt_slots_x_2window(self): # ctx=4096 / slots=8 -> per_slot_ctx=512, but 2*sliding=1024. - # SWA cells should clamp at per_slot_ctx (512), not 2*sliding. + # SWA cells clamp at per_slot_ctx (512), not 2*sliding. b = self._swa_backend() ctx = 4096 per_slot_ctx_at_8 = ctx // 8 @@ -1769,34 +1686,29 @@ class TestParallelSWAScaling: assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8) == expected def test_swa_full_does_not_scale_under_parallel(self): - # swa_full forces every layer to n_ctx; result is the all-global - # GQA-style total, which is constant in parallel. + # swa_full forces every layer to n_ctx -> all-global GQA-style + # total, constant in parallel. b = self._swa_backend() ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) for slots in (1, 2, 4, 8): assert ( - b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) - == baseline + b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline ) # ── kv_unified: no-op for memory math ────────────────────────── def test_kv_unified_is_no_op_for_memory_math(self): - # Both unified=True and unified=False must produce the same - # total bytes for every backend type and every parallel value. + # unified=True and unified=False must give the same total bytes + # for every backend type and parallel value. backends = [ ("gqa", self._gqa_backend()), ("swa", self._swa_backend()), ] for label, b in backends: for slots in (1, 2, 4, 8): - u = b._estimate_kv_cache_bytes( - 8192, "f16", n_parallel = slots, kv_unified = True - ) - nu = b._estimate_kv_cache_bytes( - 8192, "f16", n_parallel = slots, kv_unified = False - ) + u = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True) + nu = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False) assert u == nu, f"{label} parallel={slots} unified-mismatch" # ── Empirical Gemma-3 270m formula ───────────────────────────── @@ -1816,9 +1728,8 @@ class TestParallelSWAScaling: b._kv_key_length = 256 b._kv_value_length = 256 b._sliding_window = 512 - # 5-period [swa,swa,swa,swa,full] * 3 + [swa,swa,swa]: mirrors the - # bootstrap-resolved pattern for gemma3 (period 6) on an 18-layer - # model (15 SWA, 3 global). + # Mirrors the bootstrap-resolved gemma3 pattern (period 6) on an + # 18-layer model: 15 SWA, 3 global. b._sliding_window_pattern = [(i + 1) % 6 != 0 for i in range(18)] n_global = 3 n_swa = 15 @@ -1832,20 +1743,18 @@ class TestParallelSWAScaling: ), f"slots={slots}: got {got_mib} MiB, expected {expected_mib} MiB" -# --------------------------------------------------------------------------- # J3. shared_kv_layers (Gemma 3n / Gemma 4) -# --------------------------------------------------------------------------- class TestSharedKVLayers: """``.attention.shared_kv_layers`` reduces the layer count that - actually allocates KV. The trailing ``shared_kv_layers`` blocks reuse - earlier caches (Gemma 3n: 35 layers, 15 shared -> 20 allocate; Gemma 4 - same field). Unset on every other arch -> no behavioural change.""" + allocates KV. The trailing ``shared_kv_layers`` blocks reuse earlier + caches (Gemma 3n: 35 layers, 15 shared -> 20 allocate; Gemma 4 same + field). Unset on every other arch -> no behavioural change.""" def _gemma3n_backend(self, **overrides): - # Mirrors google/gemma-3n-E4B-it: 35 layers, 15 shared, - # SWA window 1024, period 5 (4 sliding + 1 full repeating). + # Mirrors google/gemma-3n-E4B-it: 35 layers, 15 shared, SWA window + # 1024, period 5 (4 sliding + 1 full repeating). defaults = { "_n_layers": 35, "_n_kv_heads": 4, @@ -1928,18 +1837,15 @@ class TestSharedKVLayers: def test_path3_pattern_loops_only_unshared_layers(self): b = self._gemma3n_backend() ctx = 8192 - # First 20 layers contribute; layers 20..34 are skipped. - # Pattern: [s,s,s,s,F] repeated. In layers 0..19: - # sliding: 16, full: 4 + # First 20 layers contribute; layers 20..34 skipped. Pattern + # [s,s,s,s,F] repeated -> in layers 0..19: sliding 16, full 4. sliding_in_unshared = sum(b._sliding_window_pattern[:20]) full_in_unshared = 20 - sliding_in_unshared assert sliding_in_unshared == 16 assert full_in_unshared == 4 kv_per = 4 * (256 + 256) * 2 swa_cells = min(ctx, 2 * 1024) - expected = ( - full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per - ) + expected = full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_shared_layers_reduces_estimate(self): @@ -1947,7 +1853,7 @@ class TestSharedKVLayers: with_shared = b._estimate_kv_cache_bytes(8192, "f16") b._shared_kv_layers = 0 without_shared = b._estimate_kv_cache_bytes(8192, "f16") - # 20/35 = 0.571 of the work; expect ~43% reduction. + # 20/35 = 0.571 of the work; ~43% reduction. ratio = with_shared / without_shared assert 0.5 < ratio < 0.65 @@ -1955,8 +1861,8 @@ class TestSharedKVLayers: b = self._gemma3n_backend() ctx = 8192 flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) - # Every unshared layer caches n_ctx; equals path-4-style sizing - # over only the 20 unshared layers. + # Every unshared layer caches n_ctx -> path-4-style sizing over + # only the 20 unshared layers. kv_per = 4 * (256 + 256) * 2 assert flagged == 20 * ctx * kv_per @@ -1974,15 +1880,15 @@ class TestSharedKVLayers: assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_shared_floors_at_one_layer(self): - # Pathological: shared >= n_layers should not zero out the cache. + # Pathological: shared >= n_layers must not zero out the cache. b = self._gqa_backend(_shared_kv_layers = 99) ctx = 4096 kv_per = 8 * (128 + 128) * 2 assert b._estimate_kv_cache_bytes(ctx, "f16") == 1 * ctx * kv_per def test_composes_with_n_parallel(self): - # Only the SWA portion of the unshared layers scales by n_parallel; - # the global portion stays constant. + # Only the SWA portion of unshared layers scales by n_parallel; + # the global portion is constant. b = self._gemma3n_backend() ctx = 8192 swa = b._sliding_window @@ -1995,9 +1901,7 @@ class TestSharedKVLayers: per_slot_ctx = max(1, ctx // slots) swa_cells = min(ctx, 2 * swa, per_slot_ctx) swa_bytes_per_slot = sliding_in_unshared * swa_cells * per_token - flagged = b._estimate_kv_cache_bytes( - ctx, "f16", n_parallel = slots, kv_unified = False - ) + flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False) assert flagged == global_bytes + slots * swa_bytes_per_slot def test_composes_with_ctx_checkpoints(self): @@ -2005,7 +1909,7 @@ class TestSharedKVLayers: ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16") with_cp = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4) - # Checkpoints only count over UNSHARED SWA layers (16 of them). + # Checkpoints count only over UNSHARED SWA layers (16 of them). sliding_in_unshared = sum(b._sliding_window_pattern[:20]) per_cp_layer = 4 * 1024 * 4 * (256 + 256) * 2 # cps * swa * heads * (k+v) * bpe assert with_cp == baseline + sliding_in_unshared * per_cp_layer @@ -2017,9 +1921,7 @@ class TestSharedKVLayers: assert b._shared_kv_layers is None -# --------------------------------------------------------------------------- # K. Lifecycle Tests -# --------------------------------------------------------------------------- class TestLifecycle: @@ -2076,7 +1978,7 @@ class TestLifecycle: assert b._n_kv_heads_by_layer is None def test_end_to_end_synthetic_mla(self): - """Full round-trip: write GGUF -> parse -> estimate.""" + """Round-trip: write GGUF -> parse -> estimate.""" b = _backend_from_gguf( "deepseek2", { @@ -2134,8 +2036,8 @@ class TestLifecycle: ) assert b._can_estimate_kv() result = b._estimate_kv_cache_bytes(131072, "f16") - # gemma3 -> period 6 from the bootstrap table, SWA cache - # double-buffered to 2 * sliding_window cells. + # gemma3 -> period 6 from bootstrap; SWA cache double-buffered to + # 2 * sliding_window cells. period = 6 kv_per = 16 * 256 * 2 expected = 0 @@ -2163,13 +2065,13 @@ class TestLifecycle: ) assert b._can_estimate_kv() assert b._shared_kv_layers == 15 - # Bootstrap table for gemma3n_text -> period 5; the resolver - # synthesises a 35-entry bool array. The first 20 entries - # (n_layers - shared) are the only ones that allocate KV. + # Bootstrap for gemma3n_text -> period 5; resolver synthesises a + # 35-entry bool array. Only the first 20 (n_layers - shared) + # allocate KV. result = b._estimate_kv_cache_bytes(8192, "f16") assert result > 0 - # Sanity: setting shared back to 0 must produce a strictly larger - # estimate (more layers allocate). + # Sanity: shared back to 0 -> strictly larger estimate (more + # layers allocate). b._shared_kv_layers = 0 unshared = b._estimate_kv_cache_bytes(8192, "f16") assert unshared > result diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py new file mode 100644 index 0000000000..7c6c514b8f --- /dev/null +++ b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py @@ -0,0 +1,523 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets. + +Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade +GitHub API are stubbed out so the suite runs without internet access and is +not subject to rate limits. +""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +_studio = Path(__file__).resolve().parent.parent.parent +if str(_studio) not in sys.path: + sys.path.insert(0, str(_studio)) + +_mod = importlib.import_module("install_llama_prebuilt") +HostInfo = _mod.HostInfo +resolve_lemonade_rocm_choice = getattr(_mod, "resolve_lemonade_rocm_choice", None) +_LEMONADE_GFX_FAMILIES = getattr(_mod, "_LEMONADE_GFX_FAMILIES", None) + +if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None: + pytest.skip("PR symbols not present - check branch", allow_module_level = True) + + +@pytest.fixture(autouse = True) +def _clear_lemonade_release_cache(): + """Prevent cross-test pollution of the lemonade release lru_cache and + selection-log dedup set when tests vary the fetch_json mock return value.""" + _cache = getattr(_mod, "_fetch_lemonade_release_cached", None) + _logged: set | None = getattr(_mod, "_lemonade_selection_logged", None) + if _cache is not None and hasattr(_cache, "cache_clear"): + _cache.cache_clear() + if _logged is not None: + _logged.clear() + yield + if _cache is not None and hasattr(_cache, "cache_clear"): + _cache.cache_clear() + if _logged is not None: + _logged.clear() + + +_STUB_TAG = "b1262" +_STUB_OS_PREFIXES = ("ubuntu", "windows") +_STUB_FAMILIES = ("gfx1151", "gfx1150", "gfx120X", "gfx110X", "gfx103X") + + +def _stub_lemonade_release() -> dict: + """Minimal lemonade release payload covering all supported GPU/OS combinations.""" + assets = [ + { + "name": f"llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip", + "browser_download_url": ( + f"https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/" + f"{_STUB_TAG}/llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip" + ), + } + for prefix in _STUB_OS_PREFIXES + for family in _STUB_FAMILIES + ] + return {"tag_name": _STUB_TAG, "assets": assets} + + +def _make_rocm_host(gfx_target: str, *, windows: bool = False) -> HostInfo: + return HostInfo( + system = "Windows" if windows else "Linux", + machine = "amd64" if windows else "x86_64", + is_windows = windows, + is_linux = not windows, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = True, + rocm_gfx_target = gfx_target, + ) + + +def _lookup_family(gfx: str) -> str | None: + for prefix, family in _LEMONADE_GFX_FAMILIES: + if gfx.startswith(prefix): + return family + return None + + +# --------------------------------------------------------------------------- +# GPU family mapping +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "gfx,expected_family", + [ + ("gfx1151", "gfx1151"), + ("gfx1150", "gfx1150"), + ("gfx1201", "gfx120X"), + ("gfx1200", "gfx120X"), + ("gfx1100", "gfx110X"), + ("gfx1030", "gfx103X"), + ], +) +def test_gpu_family_mapping(gfx, expected_family): + assert _lookup_family(gfx) == expected_family + + +def test_unknown_gpu_not_in_families(): + assert _lookup_family("gfx999") is None + + +# --------------------------------------------------------------------------- +# Asset resolution - hits real lemonade GitHub API +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "gfx,os_prefix,windows", + [ + ("gfx1151", "ubuntu", False), + ("gfx1150", "ubuntu", False), + ("gfx1201", "ubuntu", False), + ("gfx1100", "ubuntu", False), + ("gfx1030", "ubuntu", False), + ("gfx1151", "windows", True), + ("gfx1100", "windows", True), + ], +) +def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows): + host = _make_rocm_host(gfx, windows = windows) + with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): + result = resolve_lemonade_rocm_choice(host, os_prefix, "default", llama_tag = "latest") + assert result is not None, f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})" + assert _lookup_family(gfx) in result.name + assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm") + + +def test_unknown_gpu_falls_through_to_upstream(): + host = _make_rocm_host("gfx999") + result = resolve_lemonade_rocm_choice(host, "ubuntu", "default", llama_tag = "latest") + assert result is None + + +# --------------------------------------------------------------------------- +# The Linux attempt builder must plan a lemonade ROCm attempt for AMD-only hosts. +# This is the path setup.sh actually invokes (fork hosts now select from the +# manifest), so the lemonade integration is useless if it isn't wired in here. +# --------------------------------------------------------------------------- + +_linux_published_attempts = getattr(_mod, "_linux_published_attempts", None) +direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None) + +PublishedLlamaArtifact = _mod.PublishedLlamaArtifact +PublishedReleaseBundle = _mod.PublishedReleaseBundle + + +def _rocm_bundle(gfx_family: str, mapped_targets: list[str]) -> "PublishedReleaseBundle": + """A fork manifest bundle exposing a per-gfx linux-rocm artifact, so + published_rocm_choice_for_host can match the host before the lemonade + fallback is appended.""" + asset_name = f"app-b9457-linux-x64-rocm-{gfx_family}.tar.gz" + artifact = PublishedLlamaArtifact( + asset_name = asset_name, + install_kind = "linux-rocm", + runtime_line = None, + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = None, + rank = 1000, + gfx_target = gfx_family, + mapped_targets = mapped_targets, + ) + return PublishedReleaseBundle( + repo = "unslothai/llama.cpp", + release_tag = "v1.0", + upstream_tag = "b9457", + assets = {asset_name: f"https://example.invalid/{asset_name}"}, + artifacts = [artifact], + ) + + +@pytest.mark.skipif( + _linux_published_attempts is None, + reason = "Linux attempt builder not present on this branch", +) +def test_linux_attempts_include_fork_rocm_and_lemonade_for_rocm_host(): + host = _make_rocm_host("gfx1151") + bundle = _rocm_bundle("gfx1151", ["gfx1151"]) + with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): + attempts = _linux_published_attempts(host, bundle, "latest") + kinds = [a.install_kind for a in attempts] + assert "linux-rocm" in kinds, f"builder did not include any linux-rocm attempt; got {kinds}" + sources = {a.source_label for a in attempts if a.install_kind == "linux-rocm"} + # The fork's own per-gfx bundle is preferred, with the lemonade prebuilt as + # the fallback -- both must be present for a covered ROCm host. + assert "published" in sources, f"fork ROCm bundle missing; got {sources}" + assert "lemonade" in sources, f"lemonade ROCm fallback missing; got {sources}" + lemonade_attempt = next(a for a in attempts if a.source_label == "lemonade") + assert "gfx1151" in lemonade_attempt.name + + +@pytest.mark.skipif( + direct_upstream_release_plan is None, + reason = "direct release planners not present on this branch", +) +def test_direct_upstream_plan_includes_lemonade_for_windows_hip_host(): + host = _make_rocm_host("gfx1151", windows = True) + release = { + "tag_name": "b9022", + "name": "b9022", + "assets": [], + } + with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): + plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") + assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt" + kinds = [a.install_kind for a in plan.attempts] + assert "windows-hip" in kinds, f"planner did not include a lemonade HIP attempt; got {kinds}" + + +@pytest.mark.skipif( + direct_upstream_release_plan is None, + reason = "direct release planners not present on this branch", +) +def test_windows_hip_falls_back_to_upstream_when_lemonade_unavailable(): + """If lemonade returns None (e.g. gfx999 or transient API failure), the planner + must still include the upstream HIP asset rather than silently downgrading to CPU.""" + host = _make_rocm_host("gfx999", windows = True) + hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip" + release = { + "tag_name": "b9022", + "name": "b9022", + "assets": [ + { + "name": hip_asset, + "browser_download_url": f"https://example.invalid/{hip_asset}", + }, + ], + } + plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") + assert plan is not None + kinds = [a.install_kind for a in plan.attempts] + assert "windows-hip" in kinds, f"upstream HIP asset not included as fallback; got {kinds}" + hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip") + assert hip_attempt.source_label == "upstream" + + +# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ── + + +def test_lemonade_release_api_url_pinned_tag(): + """A pinned llama_tag must produce the /releases/tags/ URL.""" + assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262") + assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest") + assert _mod._lemonade_release_api_for("").endswith("/releases/latest") + + +def test_lemonade_release_api_url_encodes_tag(): + """Unexpected slashes / hashes in the tag must be URL-encoded so the URL + cannot be reshaped (defence in depth -- tags should already be sanitised + upstream).""" + url = _mod._lemonade_release_api_for("b1260/../latest") + assert "/releases/tags/b1260%2F..%2Flatest" in url + assert "//latest" not in url.split("/releases/tags/", 1)[1] + + +def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch): + """UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver.""" + monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1") + host = _make_rocm_host("gfx1151") + res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") + assert res is None + + +def test_lemonade_resolver_rejects_non_github_url(monkeypatch): + """If the GitHub API response somehow contained an off-host download URL, + the resolver must refuse to use it (lemonade assets are not in the + approved-hash manifest).""" + bad_release = { + "tag_name": _STUB_TAG, + "assets": [ + { + "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", + "browser_download_url": "https://attacker.invalid/llama.zip", + }, + ], + } + host = _make_rocm_host("gfx1151") + with patch.object(_mod, "fetch_json", return_value = bad_release): + res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") + assert res is None + + +def test_lemonade_resolver_rejects_http_scheme(): + assert not _mod._is_trusted_github_release_url( + "http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_accepts_github_cdn(): + # Real GitHub release CDN URLs carry the /github-production-release-asset- prefix. + assert _mod._is_trusted_github_release_url( + "https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_rejects_arbitrary_cdn_path(): + # A CDN URL without the release-asset path prefix must be rejected. + assert not _mod._is_trusted_github_release_url( + "https://objects.githubusercontent.com/abc/def", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_accepts_release_path(): + url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip" + assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm") + + +def test_lemonade_resolver_rejects_wrong_repo(): + """A github.com release URL for a different repo must be rejected.""" + assert not _mod._is_trusted_github_release_url( + "https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_rejects_empty_browser_download_url(): + """An asset entry with an empty browser_download_url must fall through.""" + release = { + "tag_name": _STUB_TAG, + "assets": [ + { + "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", + "browser_download_url": "", + }, + ], + } + host = _make_rocm_host("gfx1151") + with patch.object(_mod, "fetch_json", return_value = release): + res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") + assert res is None + + +def test_lemonade_runtime_patterns_include_hip_runtime(): + """linux-rocm overlay must use a broad lib glob to catch all bundled .so files. + + Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp, + ...) whose names change across ROCm releases. A broad ``lib*.so*`` glob + avoids having to enumerate every transitive dependency by name. + """ + from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice + + choice = AssetChoice( + repo = "lemonade-sdk/llamacpp-rocm", + tag = "b1262", + name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip", + url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip", + source_label = "lemonade", + install_kind = "linux-rocm", + ) + pats = runtime_patterns_for_choice(choice) + # The broad glob must be present so every .so in the lemonade bundle + # (including transitive deps added in future ROCm releases) gets overlaid. + assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}" + + +_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None) + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch): + """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES; + on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100.""" + # Two GPUs; rocminfo reports each token twice (as in the real tool output). + probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") + assert _pick_rocm_gfx_target(probe_out) == "gfx1100" + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch): + """CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None.""" + probe_out = "gfx1151\ngfx1100" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1") + assert _pick_rocm_gfx_target(probe_out) is None + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch): + """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must + return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the + two gfx1100 entries into one and making index 2 out of range.""" + # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU). + # Each GPU gets its own Agent section with a few token mentions. + probe_out = ( + "***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n" + "***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n" + "***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n" + ) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2") + assert _pick_rocm_gfx_target(probe_out) == "gfx1151" + + +# --------------------------------------------------------------------------- +# Fork release scan: Windows ROCm resolves lemonade by the requested tag +# --------------------------------------------------------------------------- + +_resolve_release_asset_choice = getattr(_mod, "resolve_release_asset_choice", None) +_ApprovedReleaseChecksums = getattr(_mod, "ApprovedReleaseChecksums", None) + + +@pytest.mark.skipif( + _resolve_release_asset_choice is None or _ApprovedReleaseChecksums is None, + reason = "fork release planner not present on this branch", +) +def test_fork_scan_windows_rocm_resolves_lemonade_by_requested_tag(): + """The fork release scan pins llama_tag to per-release upstream tags + (b9457, ...) that lemonade's own tag series never contains, so the + lemonade lookup must use the requested tag ("latest") instead. Pinning + lemonade to the per-release tag 404s on every scanned release and a + Windows ROCm host ends in a rate-limited fatal instead of the lemonade + prebuilt.""" + host = _make_rocm_host("gfx1151", windows = True) + # No windows-rocm artifact in the bundle, matching current fork releases. + bundle = _rocm_bundle("gfx1151", ["gfx1151"]) + checksums = _ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "v1.0", + upstream_tag = "b9457", + artifacts = {}, + ) + seen_urls: list[str] = [] + + def _fake_fetch(api_url, *args, **kwargs): + seen_urls.append(api_url) + if "lemonade-sdk" in api_url: + if api_url.endswith("/releases/latest"): + return _stub_lemonade_release() + raise RuntimeError(f"unexpected pinned lemonade fetch: {api_url}") + # ggml-org asset listing for the upstream HIP/CPU filename fallbacks. + return {"tag_name": "b9457", "assets": []} + + with patch.object(_mod, "fetch_json", side_effect = _fake_fetch): + attempts = _resolve_release_asset_choice( + host, + "b9457", # concrete per-release upstream tag from the scan loop + bundle, + checksums, + requested_tag = "latest", + ) + + lemonade = [a for a in attempts if a.source_label == "lemonade"] + assert lemonade, f"lemonade attempt missing for Windows ROCm host; got {attempts}" + assert "gfx1151" in lemonade[0].name + assert any( + u.endswith("/releases/latest") for u in seen_urls + ), f"lemonade was never resolved via /releases/latest; fetches: {seen_urls}" + assert not any( + "lemonade-sdk" in u and "/releases/tags/" in u for u in seen_urls + ), f"lemonade lookup was pinned to the fork release tag: {seen_urls}" + + +@pytest.mark.skipif( + direct_upstream_release_plan is None, + reason = "direct release planners not present on this branch", +) +def test_direct_upstream_plan_includes_lemonade_for_linux_rocm_host(): + """A Linux ROCm host on the ggml-org direct path (e.g. a --published-repo + override) must plan lemonade before the CPU tarball, mirroring the Windows + branch. The lemonade planning previously lived in the removed + --simple-policy dispatcher, so without this leg such hosts silently + install the CPU build.""" + host = _make_rocm_host("gfx1151") + release = { + "tag_name": "b9022", + "name": "b9022", + "assets": [ + { + "name": "llama-b9022-bin-ubuntu-x64.tar.gz", + "browser_download_url": ( + "https://github.com/ggml-org/llama.cpp/releases/download/" + "b9022/llama-b9022-bin-ubuntu-x64.tar.gz" + ), + } + ], + } + with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): + plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") + assert plan is not None, "Linux ROCm host should produce a direct plan" + kinds = [a.install_kind for a in plan.attempts] + sources = [a.source_label for a in plan.attempts] + assert "linux-rocm" in kinds, f"lemonade ROCm attempt missing; got {kinds}" + assert sources[0] == "lemonade", f"lemonade must be the first attempt; got {sources}" + assert "gfx1151" in plan.attempts[0].name diff --git a/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py b/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py index 255c04a956..666a503dbb 100644 --- a/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py +++ b/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py @@ -1,21 +1,12 @@ # 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 the cache-aware disk-space preflight in -``LlamaCppBackend.load_model``. +"""Tests for the cache-aware disk-space preflight in ``LlamaCppBackend.load_model``. -The preflight used to compare the repo's total GGUF download size against -free disk without accounting for bytes already present in the Hugging -Face cache. That made re-loading a cached large model (e.g. -``unsloth/MiniMax-M2.7-GGUF`` at 131 GB) fail cold whenever free disk was -below the full weight footprint, even though nothing needed -downloading. - -These tests exercise the preflight arithmetic in isolation by driving -``get_paths_info`` and ``try_to_load_from_cache`` through ``mock.patch``. -No network, GPU, or subprocess use. - -Cross-platform: Linux, macOS, Windows, WSL. +The preflight used to compare the repo's total GGUF size against free disk +without counting bytes already in the HF cache, so re-loading a cached large +model failed cold even though nothing needed downloading. These tests exercise +the preflight arithmetic in isolation (no network/GPU/subprocess). """ from __future__ import annotations @@ -28,10 +19,8 @@ from unittest.mock import patch import pytest -# --------------------------------------------------------------------------- -# Stub heavy / unavailable external dependencies before importing the -# module under test. Same pattern as test_kv_cache_estimation.py. -# --------------------------------------------------------------------------- +# Stub heavy / unavailable deps before importing the module under test. +# Same pattern as test_kv_cache_estimation.py. _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: @@ -99,12 +88,11 @@ def _preflight( hf_repo = "unsloth/Example-GGUF", hf_token = None, ): - """Run the preflight arithmetic as written in llama_cpp.py and return - the decision outcome as a dict. + """Run the llama_cpp.py preflight arithmetic; return the decision as a dict. ``repo_files``: list of (filename, remote_bytes). - ``cached_files``: dict {filename: on_disk_bytes} for files already in cache. - ``free_bytes``: value returned by shutil.disk_usage(cache_dir).free. + ``cached_files``: {filename: on_disk_bytes} for files already cached. + ``free_bytes``: shutil.disk_usage(cache_dir).free. """ import os import shutil @@ -112,22 +100,20 @@ def _preflight( path_infos = [_FakePathInfo(name, size) for name, size in repo_files] with tempfile.TemporaryDirectory() as tmp: - # Create SPARSE files for the cached ones so os.path.exists / - # os.path.getsize pass without actually allocating bytes on disk. - # This is critical when simulating multi-GB models. + # Sparse files so exists/getsize pass without allocating bytes on disk + # (critical for multi-GB models). cache_paths = {} for name, sz in cached_files.items(): p = Path(tmp) / name.replace("/", "_") with open(p, "wb") as fh: if sz > 0: - fh.truncate(sz) # sparse allocation: no data blocks written + fh.truncate(sz) # sparse: no data blocks written cache_paths[name] = str(p) def fake_try_to_load_from_cache(repo_id, filename): return cache_paths.get(filename) - # Mirror the same variable names and control flow as the real code - # so behavioral drift is caught immediately. + # Mirror the real code's names and control flow so drift is caught. total_bytes = sum((p.size or 0) for p in path_infos) already_cached_bytes = 0 for p in path_infos: @@ -189,8 +175,8 @@ class TestCacheAwarePreflight: assert out["would_raise_disk_error"] is False def test_partial_cache_insufficient_disk_for_rest_still_raises(self): - """Two of four shards cached; remaining 70 GB still bigger than - free disk -> preflight correctly wants to raise.""" + """Two of four shards cached; remaining 70 GB still exceeds free + disk -> preflight correctly wants to raise.""" shards = [(f"UD-Q4_K_XL/shard-{i}.gguf", 35 * GIB) for i in range(4)] cached = { shards[0][0]: shards[0][1], @@ -217,8 +203,8 @@ class TestCacheAwarePreflight: assert out["would_raise_disk_error"] is False def test_incomplete_cached_blob_is_not_credited(self): - """A partial file on disk (e.g. interrupted download) is not - counted as cached -- we still require bytes for it.""" + """A partial file on disk (e.g. interrupted download) isn't counted + as cached -- we still require bytes for it.""" shards = [("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)] partial = {"UD-Q4_K_XL/shard-0.gguf": 10 * GIB} out = _preflight( @@ -231,7 +217,7 @@ class TestCacheAwarePreflight: assert out["would_raise_disk_error"] is False def test_zero_size_path_infos_do_not_crash(self): - """A path_info with size=0 should not be credited or break the + """A path_info with size=0 must not be credited or break the arithmetic.""" shards = [("mmproj.gguf", 0), ("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)] out = _preflight( diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index 1ea76edd15..58226f938c 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -5,26 +5,14 @@ Guards two regressions in ``LlamaCppBackend.load_model``: -1. **Auto mode on weights-exceed-VRAM** (``n_ctx == 0``): when the model - weights alone exceed 90% of every GPU subset's free memory, the - auto-pick loop used to exit without matching, leaving - ``effective_ctx`` at the model's native context (e.g. 196608 for - MiniMax-M2.7). The intended default per Studio's UI spec is 4096 so - the slider lands on a usable value; the user can still drag higher - and trigger ``--fit on`` with a warning. +1. Auto mode (``n_ctx == 0``) when weights exceed every GPU subset's free + memory: auto-pick should fall back to 4096 (a usable slider value) rather + than leaving native ctx. User can still drag higher onto ``--fit on``. +2. Explicit ctx must never be silently shrunk: when KV overflows fittable + weights, honor the explicit ctx with ``--fit on`` flexing ``-ngl``. -2. **Explicit ctx silently shrunk when KV overflows**: with fittable - weights but a requested ctx whose KV cache pushes total memory over - 90% of VRAM, the old code binary-searched a smaller ctx and emitted - ``-c -ngl -1`` without informing the caller. The UI had - already surfaced its "might be slower" warning and expects the user's - explicit ctx to be honored with ``--fit on`` flexing ``-ngl`` instead. - -Tests avoid GPU probing, subprocess spawning, and GGUF I/O by driving the -post-metadata decision block directly against a stubbed instance. - -Requires no GPU, network, or external libraries beyond pytest. -Cross-platform: Linux, macOS, Windows, WSL. +Drives the post-metadata decision block against a stubbed instance: no GPU, +network, subprocess, or GGUF I/O. Cross-platform. """ from __future__ import annotations @@ -36,24 +24,20 @@ from pathlib import Path import pytest # --------------------------------------------------------------------------- -# Stub heavy / unavailable external dependencies before importing the -# module under test. Same pattern as test_kv_cache_estimation.py. +# Stub heavy/unavailable deps before importing the module under test. # --------------------------------------------------------------------------- _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# loggers _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) -# structlog _structlog_stub = _types.ModuleType("structlog") sys.modules.setdefault("structlog", _structlog_stub) -# httpx _httpx_stub = _types.ModuleType("httpx") for _exc_name in ( "ConnectError", @@ -84,6 +68,7 @@ _httpx_stub.Client = type( sys.modules.setdefault("httpx", _httpx_stub) from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx # --------------------------------------------------------------------------- @@ -102,8 +87,7 @@ def _make_backend( kv_key_length = 128, kv_value_length = 128, ): - """Create a LlamaCppBackend instance with GGUF metadata fields set and - the helpers used by the decision block stubbed out.""" + """LlamaCppBackend with GGUF metadata set and decision helpers stubbed.""" inst = LlamaCppBackend.__new__(LlamaCppBackend) inst._context_length = native_ctx inst._n_layers = n_layers @@ -131,28 +115,36 @@ def _drive( native_ctx = 131072, kv_per_token_bytes = 325_000, can_estimate_kv = True, + extra_args = None, ): """Drive the post-metadata portion of load_model with stubbed inputs. - Mirrors the decision block at llama_cpp.py:1137-1296 so we can assert - the command that would be built, without subprocesses or GPU probes. + Mirrors llama_cpp.py:1137-1296 to assert the built command, without + subprocesses or GPU probes. """ inst = _make_backend(native_ctx = native_ctx) model_size = int(model_gib * GIB) cache_type_kv = None - def fake_estimate(n_ctx_, _type = None, **_kwargs): + def fake_estimate( + n_ctx_, + _type = None, + **_kwargs, + ): return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes inst._estimate_kv_cache_bytes = fake_estimate inst._can_estimate_kv = lambda: can_estimate_kv context_length = inst._context_length + # Use the production helper, not a reimplementation, to avoid testing our own logic. + ctx_override = parse_ctx_override(extra_args) + requested_ctx = resolve_requested_ctx(extra_args, n_ctx) - effective_ctx = n_ctx if n_ctx > 0 else (context_length or 0) + effective_ctx = requested_ctx if requested_ctx > 0 else (context_length or 0) max_available_ctx = context_length or effective_ctx - if n_ctx > 0: - effective_ctx = n_ctx + if requested_ctx > 0: + effective_ctx = requested_ctx elif context_length is not None: effective_ctx = context_length else: @@ -161,7 +153,7 @@ def _drive( max_available_ctx = context_length or effective_ctx gpu_indices, use_fit = None, True - explicit_ctx = n_ctx > 0 + explicit_ctx = requested_ctx > 0 if gpus and inst._can_estimate_kv() and effective_ctx > 0: native_ctx_for_cap = context_length or effective_ctx @@ -226,9 +218,7 @@ def _drive( elif gpus: gpu_indices, use_fit = inst._select_gpus(model_size, gpus) if use_fit and not explicit_ctx: - effective_ctx = ( - min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX - ) + effective_ctx = min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX return { "c_arg": effective_ctx if effective_ctx > 0 else 0, @@ -236,6 +226,7 @@ def _drive( "gpu_indices": gpu_indices, "max_available_ctx": max_available_ctx, "original_ctx": original_ctx, + "ctx_override": ctx_override, } @@ -257,8 +248,8 @@ class TestAutoModeWeightsExceedVRAM: assert plan["c_arg"] == FALLBACK_CTX assert plan["use_fit"] is True assert plan["gpu_indices"] is None - # UI slider ceiling stays at native: user can still drag higher - # and get the "might be slower" path. + # UI slider ceiling stays at native: user can drag higher and get + # the "might be slower" path. assert plan["max_available_ctx"] == 196608 def test_multi_gpu_all_subsets_fail(self): @@ -294,9 +285,8 @@ class TestExplicitCtxRespectsUser: """``n_ctx > 0`` must never be silently shrunk.""" def test_fittable_weights_oversized_kv(self): - # 8 GB weights + 131k ctx KV on 24 GB VRAM. - # Budget = 21.6 GB, KV at 131k >> 13.6 GB remaining, so - # _select_gpus flips use_fit=True. + # 8 GB weights + 131k ctx KV on 24 GB VRAM. Budget = 21.6 GB, KV + # at 131k >> 13.6 GB remaining, so _select_gpus flips use_fit=True. plan = _drive( n_ctx = 131072, model_gib = 8, @@ -340,7 +330,7 @@ class TestExplicitCtxRespectsUser: assert plan["use_fit"] is True def test_explicit_below_floor_honored(self): - # 2048 is below --fit-ctx default; still honored since user set it. + # 2048 is below --fit-ctx default; honored since user set it. plan = _drive( n_ctx = 2048, model_gib = 8, @@ -349,6 +339,48 @@ class TestExplicitCtxRespectsUser: assert plan["c_arg"] == 2048 +# --------------------------------------------------------------------------- +# Pass-through --ctx-size participates in context fit (#5676). +# --------------------------------------------------------------------------- + + +class TestExtraArgsCtxOverride: + def test_ctx_size_extra_honored_over_auto(self): + plan = _drive( + n_ctx = 0, + model_gib = 131, + gpus = [(0, 97_000)], + native_ctx = 196608, + extra_args = ["--ctx-size", "128000"], + ) + assert plan["ctx_override"] == 128000 + assert plan["original_ctx"] == 128000 + assert plan["c_arg"] == 128000 + assert plan["use_fit"] is True + + def test_ctx_size_short_alias_honored_over_auto(self): + plan = _drive( + n_ctx = 0, + model_gib = 131, + gpus = [(0, 97_000)], + native_ctx = 196608, + extra_args = ["-c", "128000"], + ) + assert plan["c_arg"] == 128000 + assert plan["use_fit"] is True + + def test_ctx_size_extra_wins_over_first_class_field(self): + plan = _drive( + n_ctx = 4096, + model_gib = 8, + gpus = [(0, 24_000)], + native_ctx = 131072, + extra_args = ["--ctx-size", "128000"], + ) + assert plan["original_ctx"] == 128000 + assert plan["c_arg"] == 128000 + + # --------------------------------------------------------------------------- # Non-regression: fittable + auto still auto-picks largest fitting ctx # --------------------------------------------------------------------------- @@ -399,8 +431,8 @@ class TestTightFitPinsToGPU: """Models that fit at 91-95% of free VRAM must use the GPU.""" def test_rtx_4090_qwen_24gb_class(self): - # noahterbest's #5106 log: 20.8 GB model on 22805 MiB free - # GPU, ctx=4096 -> ~94% utilization, ~1.4 GiB headroom. + # noahterbest's #5106 log: 20.8 GB model on 22805 MiB free GPU, + # ctx=4096 -> ~94% utilization, ~1.4 GiB headroom. plan = _drive( n_ctx = 0, model_gib = 20.8, @@ -443,9 +475,8 @@ class TestTightFitPinsToGPU: @pytest.mark.parametrize("platform_tag", ["linux", "windows", "mac", "rocm"]) def test_identical_decision_across_platforms(platform_tag): - """The decision function takes ``[(gpu_idx, free_mib), ...]`` regardless - of how upstream (nvidia-smi / nvidia-smi.exe / Metal / rocm-smi) produced - it. Identical inputs must yield identical plans.""" + """Decision takes ``[(gpu_idx, free_mib), ...]`` regardless of source; + identical inputs must yield identical plans.""" plan_a = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)]) plan_b = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)]) assert plan_a == plan_b, platform_tag @@ -473,8 +504,8 @@ class TestClassifyGpuOffload: assert inst._classify_gpu_offload(True, [(0, 22805)]) is True def test_cpu_only_buffer_returns_false(self): - # llama-server printed buffer lines but only CPU buffers -- - # this is the silent CPU fallback symptom we want to catch. + # Buffer lines printed but only CPU buffers -- the silent CPU + # fallback symptom we want to catch. inst = self._backend( [ "load_tensors: CPU_Mapped model buffer size = 21000.0 MiB", @@ -503,8 +534,7 @@ class TestClassifyGpuOffload: assert inst._classify_gpu_offload(False, []) is None def test_user_did_not_intend_gpu_returns_none(self): - # Studio called start_llama_server without expecting GPU use; - # don't warn. + # Studio called start_llama_server without expecting GPU; don't warn. inst = self._backend( [ "load_tensors: CPU_Mapped model buffer size = 21000.0 MiB", diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index b32aeefcdb..f8e4619ded 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -3,8 +3,8 @@ """Tests for the llama.cpp prebuilt freshness check. -Pins the marker parser, the disk+memory cache, the stale decision -matrix, and fail-open behaviour on missing data. +Pins the marker parser, disk+memory cache, stale-decision matrix, and +fail-open behaviour on missing data. """ from __future__ import annotations @@ -21,12 +21,25 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) + +class _NoopLogger: + """structlog-style logger: every method swallows positional + kwargs. + + A stdlib logging.Logger rejects structlog's keyword fields (e.g. + ``logger.warning(msg, error=...)``), which leaked into the update module's + error path and failed only when this file's stub loaded first. + """ + + def __getattr__(self, _name): + return lambda *a, **k: None + + _loggers_stub = _types.ModuleType("loggers") -_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +_loggers_stub.get_logger = lambda *a, **k: _NoopLogger() sys.modules.setdefault("loggers", _loggers_stub) _structlog_stub = _types.ModuleType("structlog") -_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +_structlog_stub.get_logger = lambda *a, **k: _NoopLogger() sys.modules.setdefault("structlog", _structlog_stub) import pytest @@ -51,13 +64,18 @@ def _write_marker(install_dir: Path, **overrides) -> Path: .replace("+00:00", "Z"), } payload.update(overrides) + # The installer always writes `tag` and `release_tag` from the same release + # (a normalized base vs the full release tag), so keep the pair consistent + # when a test overrides only `tag`. + if "tag" in overrides and "release_tag" not in overrides: + payload["release_tag"] = overrides["tag"] install_dir.mkdir(parents = True, exist_ok = True) (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(payload)) return install_dir / "UNSLOTH_PREBUILT_INFO.json" def _fake_binary(install_dir: Path, *, layout: str = "cmake") -> Path: - """Stub llama-server under one of the supported install layouts.""" + """Stub llama-server under a supported install layout.""" if layout == "cmake": bin_dir = install_dir / "build" / "bin" bin_name = "llama-server" @@ -77,7 +95,7 @@ def _fake_binary(install_dir: Path, *, layout: str = "cmake") -> Path: @pytest.fixture(autouse = True) def _reset(monkeypatch, tmp_path): - # Isolate disk cache per-test; never touch the user's real cache. + # Isolate disk cache per-test; never touch the real cache. monkeypatch.setattr(fr, "_cache_dir", lambda: tmp_path / ".freshness") fr.reset_caches() yield @@ -107,8 +125,8 @@ def test_read_install_marker_finds_root_layout(tmp_path): def test_read_install_marker_finds_windows_cmake_layout(tmp_path): - # Windows cmake puts the .exe under build/bin/Release/, so the - # marker is four levels above the binary. + # Windows cmake puts the .exe under build/bin/Release/, so the marker + # is four levels above the binary. install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b8888") bin_path = _fake_binary(install_dir, layout = "windows") @@ -119,10 +137,9 @@ def test_read_install_marker_finds_windows_cmake_layout(tmp_path): @pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"]) def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo): - # The freshness check queries whichever release repo the marker - # records, so CUDA Linux (unslothai), CPU Linux x86_64 / macOS - # (ggml-org), and ROCm source-build (unslothai upstream label) - # all surface the right "latest" tag. + # The freshness check queries whichever release repo the marker records, + # so CUDA (unslothai), CPU/macOS (ggml-org), and ROCm all get the right + # "latest" tag. install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b9000", published_repo = repo) bin_path = _fake_binary(install_dir, layout = "cmake") @@ -172,9 +189,7 @@ def test_latest_published_release_returns_none_on_network_failure(monkeypatch): assert fr.latest_published_release("unslothai/llama.cpp") is None -def test_latest_published_release_keeps_old_cache_on_transient_failure( - monkeypatch, tmp_path -): +def test_latest_published_release_keeps_old_cache_on_transient_failure(monkeypatch, tmp_path): # Disk entry older than TTL + network fail -> return cached value. cache_dir = tmp_path / ".freshness" cache_dir.mkdir() @@ -188,9 +203,7 @@ def test_latest_published_release_keeps_old_cache_on_transient_failure( # check_prebuilt_freshness end-to-end. -def test_check_prebuilt_freshness_reports_stale_when_old_and_behind( - monkeypatch, tmp_path -): +def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" _write_marker( install_dir, @@ -200,9 +213,7 @@ def test_check_prebuilt_freshness_reports_stale_when_old_and_behind( .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr( - fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" - ) + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") info = fr.check_prebuilt_freshness(str(bin_path)) assert info["has_marker"] is True assert info["stale"] is True @@ -222,9 +233,7 @@ def test_check_prebuilt_freshness_not_stale_when_tag_matches(monkeypatch, tmp_pa .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr( - fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" - ) + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") info = fr.check_prebuilt_freshness(str(bin_path)) assert info["stale"] is False assert info["installed_tag"] == "b9300" @@ -242,9 +251,7 @@ def test_check_prebuilt_freshness_not_stale_within_threshold(monkeypatch, tmp_pa .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr( - fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" - ) + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") info = fr.check_prebuilt_freshness(str(bin_path)) assert info["stale"] is False assert info["age_days"] == 1 @@ -257,9 +264,7 @@ def test_check_prebuilt_freshness_fails_open_without_marker(tmp_path): assert info["stale"] is False -def test_check_prebuilt_freshness_fails_open_when_github_unreachable( - monkeypatch, tmp_path -): +def test_check_prebuilt_freshness_fails_open_when_github_unreachable(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" _write_marker( install_dir, @@ -276,15 +281,11 @@ def test_check_prebuilt_freshness_fails_open_when_github_unreachable( assert info["latest_tag"] is None -def test_check_prebuilt_freshness_handles_unparseable_install_timestamp( - monkeypatch, tmp_path -): +def test_check_prebuilt_freshness_handles_unparseable_install_timestamp(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b9190", installed_at_utc = "not-a-date") bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr( - fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" - ) + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") info = fr.check_prebuilt_freshness(str(bin_path)) assert info["stale"] is False assert info["age_days"] is None @@ -300,9 +301,7 @@ def test_check_prebuilt_freshness_respects_custom_threshold(monkeypatch, tmp_pat .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr( - fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" - ) + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") info = fr.check_prebuilt_freshness(str(bin_path), threshold_days = 1) assert info["stale"] is True @@ -311,9 +310,7 @@ def test_check_prebuilt_freshness_respects_custom_threshold(monkeypatch, tmp_pat def test_format_stale_warning_contains_actionable_command(): - msg = fr.format_stale_warning( - {"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 5} - ) + msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 5}) assert "b9190" in msg assert "b9300" in msg assert "5 days" in msg @@ -321,8 +318,118 @@ def test_format_stale_warning_contains_actionable_command(): def test_format_stale_warning_singular_day(): - msg = fr.format_stale_warning( - {"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1} - ) + msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1}) assert "1 day" in msg assert "1 days" not in msg + + +# parse_base_build / is_behind. + + +def test_parse_base_build(): + assert fr.parse_base_build("b9596") == 9596 + assert fr.parse_base_build(" b9596 ") == 9596 + assert fr.parse_base_build("b9596-mix-e6f2453") == 9596 # mix suffix doesn't defeat it + assert fr.parse_base_build("9596") is None + assert fr.parse_base_build("master-abc") is None + assert fr.parse_base_build("") is None + assert fr.parse_base_build(None) is None + + +@pytest.mark.parametrize( + "installed, latest, expected", + [ + ( + "b9596-mix-e6f2453", + "b9596-mix-e6f2453", + False, + ), # already on the mix latest -> not behind + ("b9596", "b9594", False), # latest is an older build -> downgrade guard + ("b9596", "b9594-mix-xxx", False), # older mix latest -> still guarded + ("b9500", "b9596-mix-e6f2453", True), # newer base -> behind + ("b9596-mix-aaa", "b9596-mix-bbb", True), # new mix at same base -> behind + ("b9596", "b9596-mix-bbb", True), # clean -> mix at same base -> behind + ("b9596-mix-aaa", "b9596", False), # bare base never supersedes a mix install + ("b9596", "b9596", False), # identical -> not behind + (" b9596 ", "b9596", False), # whitespace-only diff -> not behind + ("master-abc", "master-def", True), # non-bNNNN both -> plain inequality + ("master-abc", "master-abc", False), + (None, "b9596", False), + ("b9596", None, False), + ], +) +def test_is_behind(installed, latest, expected): + assert fr.is_behind(installed, latest) is expected + + +def test_check_prebuilt_freshness_not_behind_on_mix_latest(monkeypatch, tmp_path): + # Installed the mix latest: marker base tag b9596, full release_tag with sha, + # GitHub latest is that same full tag. Must not report behind (sticky bug). + install_dir = tmp_path / "llama.cpp" + _write_marker(install_dir, tag = "b9596", release_tag = "b9596-mix-e6f2453") + bin_path = _fake_binary(install_dir, layout = "root") + monkeypatch.setattr( + fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453" + ) + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["behind"] is False + assert info["stale"] is False + + +def test_check_prebuilt_freshness_downgrade_guard(monkeypatch, tmp_path): + # A lagging latest (older build than installed) must never read as behind/stale. + install_dir = tmp_path / "llama.cpp" + _write_marker( + install_dir, + tag = "b9585", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 30)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(install_dir, layout = "root") + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["behind"] is False + assert info["stale"] is False + + +def test_fetch_latest_release_tag_uses_publish_time(monkeypatch): + # Resolves newest by published_at (like the installer), skips drafts/prereleases, + # and does NOT just take GitHub's first/`/releases/latest` item. + import urllib.request + + class _Resp: + def __init__(self, payload): + self._p = json.dumps(payload).encode() + + def read(self): + return self._p + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + payload = [ + { + "tag_name": "b9518", + "draft": False, + "prerelease": False, + "published_at": "2026-06-04T21:11:19Z", + }, + { + "tag_name": "b9596-mix-e6f2453", + "draft": False, + "prerelease": False, + "published_at": "2026-06-11T22:50:41Z", + }, + { + "tag_name": "b9999-draft", + "draft": True, + "prerelease": False, + "published_at": "2026-06-12T00:00:00Z", + }, + ] + monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = 5.0: _Resp(payload)) + assert fr._fetch_latest_release_tag("unslothai/llama.cpp") == "b9596-mix-e6f2453" diff --git a/studio/backend/tests/test_llama_cpp_load_progress.py b/studio/backend/tests/test_llama_cpp_load_progress.py index f46751b798..cc1c6256a8 100644 --- a/studio/backend/tests/test_llama_cpp_load_progress.py +++ b/studio/backend/tests/test_llama_cpp_load_progress.py @@ -3,34 +3,20 @@ """Tests for ``LlamaCppBackend.load_progress()``. -The chat settings flow and the training overlay both show a generic -"Starting model..." spinner during the window after a GGUF download -finishes and before llama-server reports healthy. For small models -that window is a second or two and nobody notices. For large MoE GGUFs -(MiniMax-M2.7, Qwen3.5-397B-A17B, etc.) the llama-server process spends -minutes in kernel state D, paging tens or hundreds of GB of shards -into the page cache. The UI has no way to show a real progress bar, -rate, or ETA during that window. +For large MoE GGUFs, llama-server spends minutes paging shards into the page +cache after download. ``load_progress()`` samples ``/proc//status VmRSS`` +against the total shard size on disk so the UI can render a real bar plus +rate/ETA. Contract pinned here: -``load_progress()`` samples ``/proc//status VmRSS`` (what the -kernel has actually paged in) against the total shard file size on -disk, so the frontend can render a real bar plus rate/ETA. This -module pins that contract: - - * returns ``None`` when no load is in flight - * returns ``{"phase": "mmap", ...}`` while the subprocess is alive - but ``_healthy`` is False - * returns ``{"phase": "ready", ...}`` once ``_healthy`` flips - * ``bytes_total`` is derived from the resolved on-disk path - (which the paired fix assigns to ``self._gguf_path`` on both the - local-GGUF and HF-download code paths) + * ``None`` when no load is in flight + * ``{"phase": "mmap", ...}`` while the subprocess is alive but ``_healthy`` is False + * ``{"phase": "ready", ...}`` once ``_healthy`` flips + * ``bytes_total`` derived from the resolved on-disk path (``self._gguf_path``) * ``bytes_loaded`` is VmRSS in bytes, capped by total, rounded - * ``fraction`` is clamped to 0..1 and rounded to 4 decimal places + * ``fraction`` clamped to 0..1, rounded to 4 dp -Linux-only via ``/proc``. On platforms without ``/proc`` the method -returns ``None`` instead of raising. -Cross-platform test: skips cleanly on macOS / Windows if ``/proc`` is -not available. +Linux-only via ``/proc``; returns ``None`` (not raises) without it, so tests +skip cleanly on macOS / Windows. """ from __future__ import annotations @@ -44,10 +30,7 @@ from unittest.mock import patch import pytest -# --------------------------------------------------------------------------- -# Stub heavy / unavailable external dependencies before importing the -# module under test. Same pattern as test_kv_cache_estimation.py. -# --------------------------------------------------------------------------- +# Stub heavy / unavailable deps before importing the module under test. _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: @@ -106,7 +89,7 @@ def _make_instance(): class _FakeProc: - """Minimal stand-in for subprocess.Popen that just carries a pid.""" + """Minimal stand-in for subprocess.Popen carrying just a pid.""" def __init__(self, pid: int): self.pid = pid @@ -150,7 +133,6 @@ class TestLoadProgressSingleShard: def fake_open(path, *args, **kwargs): if str(path).startswith("/proc/"): import io - return io.StringIO(f"Name:\ttest\nVmRSS:\t{10 * 1024 ** 2}\tkB\n") return open(path, *args, **kwargs) # fall through @@ -175,7 +157,6 @@ class TestLoadProgressSingleShard: def fake_open(path, *args, **kwargs): if str(path).startswith("/proc/"): import io - return io.StringIO(f"VmRSS:\t{8 * 1024 ** 2}\tkB\n") return open(path, *args, **kwargs) @@ -190,7 +171,7 @@ class TestLoadProgressSingleShard: class TestLoadProgressMultiShard: - """Shard-aware total: for ``*-00001-of-00004.gguf`` primaries the + """Shard-aware total: for ``*-00001-of-00004.gguf`` primaries, the method sums sibling files with the same prefix.""" def test_sharded_total_aggregates_siblings(self, tmp_path): @@ -199,7 +180,7 @@ class TestLoadProgressMultiShard: tmp_path / f"model-{i:05d}-of-00004.gguf", size_bytes = 20 * 1024**3, ) - # Drop an unrelated .gguf in the same folder -- must not be counted. + # An unrelated .gguf in the same folder -- must not be counted. _write_sparse_file(tmp_path / "mmproj-BF16.gguf", 2 * 1024**3) inst = _make_instance() @@ -210,7 +191,6 @@ class TestLoadProgressMultiShard: def fake_open(path, *args, **kwargs): if str(path).startswith("/proc/"): import io - return io.StringIO("VmRSS:\t0\tkB\n") return open(path, *args, **kwargs) @@ -233,7 +213,6 @@ class TestLoadProgressDegradation: def fake_open(path, *args, **kwargs): if str(path).startswith("/proc/"): import io - return io.StringIO("VmRSS:\t1024\tkB\n") return open(path, *args, **kwargs) diff --git a/studio/backend/tests/test_llama_cpp_load_progress_live.py b/studio/backend/tests/test_llama_cpp_load_progress_live.py index beed8713c1..98e19944dd 100644 --- a/studio/backend/tests/test_llama_cpp_load_progress_live.py +++ b/studio/backend/tests/test_llama_cpp_load_progress_live.py @@ -3,20 +3,10 @@ """Live, no-mock integration test for ``LlamaCppBackend.load_progress()``. -The companion files (``test_llama_cpp_load_progress.py`` and -``test_llama_cpp_load_progress_matrix.py``) patch ``builtins.open`` to -feed synthetic VmRSS values. This file is the opposite: it uses **real** -subprocesses, **real** file sizes, and the **real** ``/proc`` -interface. It is the sanity check that the contract we keep in the -mocked tests still maps to what the kernel actually returns on a live -Linux system. - -Why both: the mocked tests can be fooled by a buggy implementation that -parses ``/proc`` output in a format the kernel no longer uses, or that -makes assumptions about ``Path.stat()`` vs ``os.path.getsize``. This -file hits the real APIs so any format drift gets caught. - -Skipped cleanly on non-Linux (no ``/proc``). +The companion mocked tests patch ``builtins.open`` for synthetic VmRSS values; +this one uses real subprocesses, file sizes, and ``/proc`` so format drift the +mocks can't see (kernel ``/proc`` layout, stat vs getsize) gets caught. Skipped +on non-Linux (no ``/proc``). """ from __future__ import annotations @@ -30,10 +20,7 @@ from pathlib import Path import pytest -# --------------------------------------------------------------------------- -# Same stubs as the matrix file (keep self-contained so the file can be -# run standalone as well as via the full suite). -# --------------------------------------------------------------------------- +# Same stubs as the matrix file (self-contained for standalone + full-suite runs). _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: @@ -75,7 +62,11 @@ pytestmark = pytest.mark.skipif( ) -def _make_backend(pid: int, gguf_path: str, healthy: bool = False): +def _make_backend( + pid: int, + gguf_path: str, + healthy: bool = False, +): inst = LlamaCppBackend.__new__(LlamaCppBackend) inst._process = type("P", (), {"pid": pid})() inst._gguf_path = gguf_path @@ -84,8 +75,8 @@ def _make_backend(pid: int, gguf_path: str, healthy: bool = False): def test_live_rss_matches_kernel_vmrss(tmp_path): - """Spawn a real child, let it allocate real bytes, confirm - ``bytes_loaded`` tracks the kernel's VmRSS within a sane tolerance.""" + """Spawn a real child, let it allocate real bytes, confirm ``bytes_loaded`` + tracks the kernel's VmRSS within a sane tolerance.""" # Child that allocates ~100 MB of zero'd bytes and then idles. script = tmp_path / "burn.py" script.write_text( @@ -108,7 +99,7 @@ def test_live_rss_matches_kernel_vmrss(tmp_path): ready = proc.stdout.readline() assert ready.strip() == b"ready" - # Create a fake 200 MB sparse gguf so bytes_total is concrete. + # Fake 200 MB sparse gguf so bytes_total is concrete. gguf = tmp_path / "model.gguf" with open(gguf, "wb") as f: f.truncate(200 * 1024 * 1024) @@ -119,8 +110,8 @@ def test_live_rss_matches_kernel_vmrss(tmp_path): assert out is not None, "load_progress returned None for live pid" assert out["phase"] == "mmap" assert out["bytes_total"] == 200 * 1024 * 1024 - # VmRSS for the Python child includes the interpreter + the 100MB - # buffer, so a realistic floor is 50 MB and ceiling is 200 MB. + # VmRSS for the Python child includes the interpreter + 100MB buffer, + # so a realistic floor is 50 MB and ceiling is 200 MB. assert ( out["bytes_loaded"] >= 50 * 1024 * 1024 ), f"bytes_loaded unexpectedly low: {out['bytes_loaded']}" @@ -149,8 +140,8 @@ def test_live_ready_phase_when_healthy(tmp_path): def test_live_dead_pid_returns_none(tmp_path): - """A recently-dead pid may linger in /proc for ms; use a clearly - invalid id so the read reliably fails.""" + """A recently-dead pid may linger in /proc for ms; use a clearly invalid id + so the read reliably fails.""" gguf = tmp_path / "m.gguf" gguf.touch() @@ -160,8 +151,8 @@ def test_live_dead_pid_returns_none(tmp_path): def test_live_shard_aggregation_counts_real_files(tmp_path): - """With 4 real sibling shards on disk, ``bytes_total`` equals their - summed size to the byte.""" + """With 4 real sibling shards on disk, ``bytes_total`` equals their summed + size to the byte.""" shard_size = 7 * 1024 * 1024 # 7 MB each for i in range(1, 5): f = tmp_path / f"model-{i:05d}-of-00004.gguf" @@ -182,8 +173,8 @@ def test_live_shard_aggregation_counts_real_files(tmp_path): def test_live_repeated_polling_stays_sane(tmp_path): - """Sampling the same backend 20 times should not raise or produce - non-numeric output, even under normal kernel RSS jitter.""" + """Sampling the same backend 20 times must not raise or produce non-numeric + output, even under normal kernel RSS jitter.""" gguf = tmp_path / "m.gguf" with open(gguf, "wb") as f: f.truncate(500 * 1024 * 1024) diff --git a/studio/backend/tests/test_llama_cpp_load_progress_matrix.py b/studio/backend/tests/test_llama_cpp_load_progress_matrix.py index a88450ec0b..5c4d9106d8 100644 --- a/studio/backend/tests/test_llama_cpp_load_progress_matrix.py +++ b/studio/backend/tests/test_llama_cpp_load_progress_matrix.py @@ -3,29 +3,12 @@ """Extended test matrix for ``LlamaCppBackend.load_progress()``. -Companion to ``test_llama_cpp_load_progress.py`` (which pins the basic -contract). This file widens coverage to the edge cases that bit users -or were hypothesized to bite them on cross-platform installs: +Companion to ``test_llama_cpp_load_progress.py`` (basic contract). Covers +cross-platform edge cases: platform matrix (/proc absence), VmRSS parsing, +filesystem edges (HF-cache symlinks, broken/missing/relative paths), shard +aggregation, lifecycle races, concurrent sampling, and fraction bounds. - * Platform matrix — macOS/Windows simulation via ``/proc`` absence. - * ``VmRSS`` parsing — tab vs space delimiter, missing line, malformed - integer. - * Filesystem edges — HF-cache symlinks, broken symlinks, nonexistent - paths, relative paths. - * Shard aggregation — partial multi-shard downloads where some shards - are still ``.incomplete``, two shard series in the same dir, - ``mmproj-*.gguf`` sibling exclusion for non-sharded primaries, - single-file models. - * Lifecycle races — process set before ``_gguf_path`` is assigned, - process dead mid-sample, ``_healthy`` flipped to True. - * Concurrent sampling — 10 threads × 50 iterations against a single - backend, hitting real ``/proc`` (no mocks — see the note in - ``TestConcurrentSampling`` for why). - * Fraction bounds — capped at 1.0 when RSS exceeds total; 0.0 when - total is zero. - -All tests are Linux-only in practice (we stub ``/proc`` where needed). -The stable subset runs in well under a second. +Linux-only in practice (``/proc`` stubbed where needed). """ from __future__ import annotations @@ -40,10 +23,7 @@ from unittest.mock import patch import pytest -# --------------------------------------------------------------------------- -# Stub heavy / unavailable external dependencies before importing the -# module under test. Same pattern as test_llama_cpp_load_progress.py. -# --------------------------------------------------------------------------- +# Stub heavy/unavailable deps before importing the module under test. _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: @@ -113,7 +93,7 @@ def _sparse(path, size): def _fake_proc_reader(rss_kb): - """Return an ``open()`` replacement that fakes /proc reads with a VmRSS line.""" + """An ``open()`` replacement faking /proc reads with a VmRSS line.""" def fake_open(path, *args, **kwargs): if str(path).startswith("/proc/"): @@ -129,8 +109,8 @@ def _fake_proc_reader(rss_kb): class TestPlatformMatrix: - """The method is Linux-first via /proc. On macOS/Windows it must - degrade to None rather than crash.""" + """Linux-first via /proc. On macOS/Windows must degrade to None + rather than crash.""" def test_linux_live_proc_is_self_pid(self, tmp_path): """Self-pid /proc read uses the real kernel interface.""" @@ -144,7 +124,7 @@ class TestPlatformMatrix: assert out is not None assert out["phase"] == "mmap" assert out["bytes_total"] == 1 * 1024**3 - # Our Python process has some RSS -- just sanity-check positive. + # Our process has some RSS -- sanity-check it's positive. assert out["bytes_loaded"] > 0 def test_macos_no_proc_returns_none(self, tmp_path): @@ -199,7 +179,7 @@ class TestVmRSSParsing: assert out["bytes_loaded"] == 2 * 1024**3 def test_space_separated_fallback(self, tmp_path): - """Some kernels emit single-space rather than tab.""" + """Some kernels emit a single space, not a tab.""" gguf = tmp_path / "m.gguf" _sparse(gguf, 4 * 1024**3) inst = _make() @@ -235,8 +215,8 @@ class TestVmRSSParsing: assert out["fraction"] == 0.0 def test_malformed_vmrss_value(self, tmp_path): - """Non-integer VmRSS value should be treated as if the line were - absent (early ValueError caught).""" + """Non-integer VmRSS is treated like an absent line (ValueError + caught).""" gguf = tmp_path / "m.gguf" _sparse(gguf, 1 * 1024**3) inst = _make() @@ -250,7 +230,7 @@ class TestVmRSSParsing: with patch("builtins.open", side_effect = fake_open): out = inst.load_progress() - # The implementation catches ValueError on int() and returns None. + # int() ValueError is caught and returns None. assert out is None @@ -262,7 +242,7 @@ class TestVmRSSParsing: class TestFilesystemEdges: def test_symlink_primary_follows_to_blob(self, tmp_path): """HF cache stores blobs under blobs/ and symlinks them from - snapshots/. The method must follow the symlink.""" + snapshots/. Must follow the symlink.""" blob = tmp_path / "blob" _sparse(blob, 12 * 1024**3) snap = tmp_path / "snap" @@ -300,7 +280,7 @@ class TestFilesystemEdges: def test_relative_gguf_path(self, tmp_path): """Relative paths shouldn't crash; behaviour depends on CWD but - the method must not raise.""" + must not raise.""" cwd = os.getcwd() try: os.chdir(tmp_path) @@ -323,11 +303,11 @@ class TestFilesystemEdges: class TestShardAggregation: def test_partial_multi_shard_download(self, tmp_path): - """Primary present but shards 2..N still downloading as - ``.incomplete``. Sums only the fully-arrived ``.gguf`` files.""" + """Primary present but shards 2..N still ``.incomplete``. Sums + only the fully-arrived ``.gguf`` files.""" _sparse(tmp_path / "m-00001-of-00004.gguf", 30 * 1024**3) _sparse(tmp_path / "m-00002-of-00004.gguf", 30 * 1024**3) - # 3 and 4 still downloading as .incomplete + # 3 and 4 still downloading as .incomplete. _sparse(tmp_path / "m-00003-of-00004.gguf.incomplete", 5 * 1024**3) inst = _make() inst._process = _Proc(os.getpid()) @@ -337,8 +317,8 @@ class TestShardAggregation: assert out["bytes_total"] == 60 * 1024**3 # only the .gguf siblings def test_two_shard_series_in_same_dir(self, tmp_path): - """Defensive: if two quant series share a dir, prefix filter - only sums siblings of the chosen primary.""" + """Defensive: when two quant series share a dir, the prefix + filter sums only siblings of the chosen primary.""" for i in range(1, 3): _sparse(tmp_path / f"m_q4-{i:05d}-of-00002.gguf", 10 * 1024**3) _sparse(tmp_path / f"m_q8-{i:05d}-of-00002.gguf", 20 * 1024**3) @@ -351,7 +331,7 @@ class TestShardAggregation: def test_mmproj_sibling_not_counted(self, tmp_path): """Vision models drop an ``mmproj-*.gguf`` alongside. For a - single-file (non-sharded) primary we only count the primary.""" + single-file (non-sharded) primary, count only the primary.""" _sparse(tmp_path / "m.gguf", 8 * 1024**3) _sparse(tmp_path / "mmproj-BF16.gguf", 2 * 1024**3) inst = _make() @@ -359,7 +339,7 @@ class TestShardAggregation: inst._gguf_path = str(tmp_path / "m.gguf") with patch("builtins.open", side_effect = _fake_proc_reader(0)): out = inst.load_progress() - # Non-sharded primary: only the primary is counted. + # Non-sharded: only the primary is counted. assert out["bytes_total"] == 8 * 1024**3 def test_single_file_model(self, tmp_path): @@ -381,7 +361,7 @@ class TestShardAggregation: class TestLifecycleRaces: def test_process_set_but_gguf_path_not_yet(self, tmp_path): - """Moment between Popen and self._gguf_path=model_path.""" + """Window between Popen and self._gguf_path=model_path.""" inst = _make() inst._process = _Proc(os.getpid()) inst._gguf_path = None @@ -418,14 +398,10 @@ class TestLifecycleRaces: class TestConcurrentSampling: def test_parallel_invocations_never_raise(self, tmp_path): - """Many concurrent samplers hitting the same backend must not raise. + """Many concurrent samplers on one backend must not raise. - We intentionally do NOT patch ``builtins.open`` here because - ``unittest.mock.patch`` is not thread-safe: interleaved - enter/exit across threads can leak a Mock into ``builtins.open`` - and poison every subsequent test in the session. Instead, we - let each thread hit the real ``/proc/self/status`` of the test - process, which is exactly the code path that matters in prod. + No ``builtins.open`` patch: ``mock.patch`` isn't thread-safe and could + leak a Mock into ``open``. Each thread hits the real ``/proc/self/status``. """ _sparse(tmp_path / "m.gguf", 1 * 1024**3) inst = _make() diff --git a/studio/backend/tests/test_llama_cpp_max_context_threshold.py b/studio/backend/tests/test_llama_cpp_max_context_threshold.py index 22e4cda7d1..310aaf6c0f 100644 --- a/studio/backend/tests/test_llama_cpp_max_context_threshold.py +++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py @@ -3,22 +3,20 @@ """Tests for the ``max_context_length`` warning-threshold semantics. -``/api/inference/status.max_context_length`` is what the ctx slider in -the chat settings sheet reads to decide when to render the "Exceeds -estimated VRAM capacity. The model may use system RAM." warning: +The ctx slider in the chat settings sheet reads +``/api/inference/status.max_context_length`` to decide when to render the +"Exceeds estimated VRAM capacity. The model may use system RAM." warning: ctxDisplayValue > ggufMaxContextLength → show warning -For models whose weights fit on some GPU subset, the warning threshold -is the largest ctx that fits fully in VRAM (the binary-search cap from -``_fit_context_to_vram``). For models whose weights exceed 90% of every -GPU subset's free memory, the warning must fire as soon as the user -drags above the 4096 spec default (otherwise a user loading e.g. -MiniMax-M2.7 on a 97 GB GPU sees a slider up to 196608 with no -indication that any value above 4096 will trigger ``--fit on`` and -degrade performance). +When weights fit on some GPU subset, the threshold is the largest ctx that +fits fully in VRAM (the binary-search cap from ``_fit_context_to_vram``). +When weights exceed 90% of every GPU subset's free memory, the warning must +fire as soon as the user drags above the 4096 spec default (otherwise loading +e.g. MiniMax-M2.7 on a 97 GB GPU shows a slider up to 196608 with no hint that +any value above 4096 triggers ``--fit on`` and degrades performance). -These tests pin both cases. No GPU probing, no subprocess, no GGUF I/O. +These tests pin both cases. No GPU probing, subprocess, or GGUF I/O. Cross-platform: Linux, macOS, Windows, WSL. """ @@ -30,10 +28,8 @@ from pathlib import Path import pytest -# --------------------------------------------------------------------------- -# Stub heavy / unavailable external dependencies before importing the -# module under test. Same pattern as test_kv_cache_estimation.py. -# --------------------------------------------------------------------------- +# Stub heavy / unavailable deps before importing the module under test. +# Same pattern as test_kv_cache_estimation.py. _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: @@ -81,9 +77,7 @@ sys.modules.setdefault("httpx", _httpx_stub) from core.inference.llama_cpp import LlamaCppBackend -# --------------------------------------------------------------------------- # Helpers -# --------------------------------------------------------------------------- GIB = 1024**3 @@ -109,10 +103,14 @@ def _make_backend(native_ctx = 131072): return inst -def _compute_max_available_ctx(native_ctx, model_gib, gpus, kv_per_token_bytes = 325_000): - """Run the ceiling-probe block from load_model and return the final - ``max_available_ctx`` value the backend would assign to - ``_max_context_length``. +def _compute_max_available_ctx( + native_ctx, + model_gib, + gpus, + kv_per_token_bytes = 325_000, +): + """Run load_model's ceiling-probe block and return the final + ``max_available_ctx`` the backend would assign to ``_max_context_length``. """ inst = _make_backend(native_ctx = native_ctx) model_size = int(model_gib * GIB) @@ -152,14 +150,12 @@ def _compute_max_available_ctx(native_ctx, model_gib, gpus, kv_per_token_bytes = return max_available_ctx -# --------------------------------------------------------------------------- # Weights exceed every GPU subset's VRAM (MiniMax-M2.7-like) -# --------------------------------------------------------------------------- class TestMaxContextLengthForWeightsExceedVRAM: - """The UI ``max_context_length`` threshold must fall back to 4096 so - the warning fires as soon as the user drags above the spec default. + """UI ``max_context_length`` must fall back to 4096 so the warning fires + as soon as the user drags above the spec default. """ def test_minimax_like(self): @@ -181,8 +177,8 @@ class TestMaxContextLengthForWeightsExceedVRAM: assert got == 4096 def test_native_below_fallback_is_preserved(self): - """If the model's native ctx is itself smaller than 4096, do not - advertise a larger value than the model supports.""" + """If native ctx is itself below 4096, don't advertise a larger value + than the model supports.""" got = _compute_max_available_ctx( native_ctx = 2048, model_gib = 200, @@ -191,9 +187,7 @@ class TestMaxContextLengthForWeightsExceedVRAM: assert got == 2048 -# --------------------------------------------------------------------------- # Fittable models (regression guard) -# --------------------------------------------------------------------------- class TestMaxContextLengthForFittableModels: @@ -231,9 +225,7 @@ class TestMaxContextLengthForFittableModels: assert got >= 131072 - 256 # rounded to 256 boundary -# --------------------------------------------------------------------------- # Property plumbing -# --------------------------------------------------------------------------- class TestMaxContextLengthProperty: diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py new file mode 100644 index 0000000000..6ef94545fe --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -0,0 +1,187 @@ +# 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 the llama-server mmproj text-only fallback. + +A GGUF vision model is launched with ``--mmproj ``. When the +installed llama.cpp prebuilt is older than the model's projector format, +llama-server aborts at startup with ``clip.cpp:NNNN: Unknown projector +type`` (exit -6). load_model now retries once WITHOUT ``--mmproj`` so the +base model still loads text-only, warns the user to update llama.cpp, and +marks the session non-vision. These tests pin the two decision helpers: +``_is_projector_incompatibility`` (when to retry) and ``_strip_mmproj_args`` +(how the retry argv is built). Unrelated failures must NOT trigger a retry. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Match the stubbing pattern in sibling tests so the module imports in a +# lightweight env without fastapi. +_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("structlog") +sys.modules.setdefault("structlog", _structlog_stub) +if not hasattr(sys.modules["structlog"], "get_logger"): + sys.modules["structlog"].get_logger = _structlog_stub.get_logger + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + +_detect = LlamaCppBackend._is_projector_incompatibility +_strip = LlamaCppBackend._strip_mmproj_args + +# Real abort captured loading gemma-4 on a 3-day-old prebuilt (build b9496). +_GEMMA4_OLD_LLAMACPP_OUT = ( + "srv load_model: loading model 'gemma-4-E2B-it-UD-Q4_K_XL.gguf'\n" + "/build_work/src/llama.cpp-b9496/tools/mtmd/clip.cpp:4391: " + "Unknown projector type\n" + "libggml-base.so.0(ggml_abort+0x152)\n" + "libmtmd.so.0(clip_n_mmproj_embd)\n" +) +# Unrelated failures that must keep their own handling (no projector retry). +_OOM_OUT = ( + "ggml_backend_cuda_buffer_type_alloc_buffer: allocating 12000.00 MiB on " + "device 0: cudaMalloc failed: out of memory" +) +_BAD_ARCH_OUT = "llama_model_load: error loading model: unknown model architecture: 'qwen_image'" +_PORT_OUT = "srv start: failed to bind: address already in use" +_MISSING_OUT = "error: failed to open GGUF file: no such file or directory" +# A healthy startup log that merely mentions the projector must not match. +_HEALTHY_VISION_OUT = ( + "Using mmproj for vision: /cache/mmproj-F16.gguf\n" + "clip_model_loader: loaded meta data with 20 key-value pairs\n" + "srv update_slots: all slots are idle" +) + + +class TestProjectorIncompatibilityDetector: + def test_gemma4_on_old_llamacpp_triggers_retry(self): + # Headline case: a 3-day-old llama.cpp aborts on Gemma-4's projector. + assert _detect(_GEMMA4_OLD_LLAMACPP_OUT) is True + + @pytest.mark.parametrize( + "out", + [ + "clip.cpp:4391: Unknown projector type", + "error: unsupported projector type for this model", + "llama_mmproj: unsupported mmproj file version", + "clip.cpp: projector type 'gemma4' is not supported", + ], + ) + def test_projector_format_errors_match(self, out): + assert _detect(out) is True + + def test_case_insensitive(self): + assert _detect("UNKNOWN PROJECTOR TYPE") is True + + @pytest.mark.parametrize( + "out", + [ + _OOM_OUT, + _BAD_ARCH_OUT, + _PORT_OUT, + _MISSING_OUT, + _HEALTHY_VISION_OUT, + "", + # bare multimodal words without a failure term must not match + "loading clip model", + "mmproj file resolved from cache", + ], + ) + def test_unrelated_failures_do_not_retry(self, out): + assert _detect(out) is False + + +# A realistic vision launch argv (mirrors the live "Starting llama-server" +# command), projector pair at the end. +_VISION_CMD = [ + "/home/u/.unsloth/llama.cpp/build/bin/llama-server", + "-m", + "/cache/gemma-4-E2B-it-UD-Q4_K_XL.gguf", + "--port", + "55473", + "-c", + "131072", + "--parallel", + "1", + "--flash-attn", + "on", + "--no-context-shift", + "-ngl", + "-1", + "--threads", + "-1", + "--jinja", + "--spec-default", + "--mmproj", + "/cache/mmproj-F16.gguf", +] + + +class TestStripMmprojArgs: + def test_removes_mmproj_pair(self): + stripped = _strip(_VISION_CMD) + assert "--mmproj" not in stripped + assert "/cache/mmproj-F16.gguf" not in stripped + + def test_preserves_every_text_flag(self): + stripped = _strip(_VISION_CMD) + for flag in ( + "-m", + "/cache/gemma-4-E2B-it-UD-Q4_K_XL.gguf", + "--port", + "55473", + "-c", + "131072", + "-ngl", + "-1", + "--jinja", + "--spec-default", + "--flash-attn", + "on", + ): + assert flag in stripped + # Exactly the two projector tokens are dropped. + assert len(stripped) == len(_VISION_CMD) - 2 + + def test_strips_mmproj_in_the_middle(self): + cmd = ["llama-server", "--mmproj", "/p/mm.gguf", "-c", "4096", "--jinja"] + assert _strip(cmd) == ["llama-server", "-c", "4096", "--jinja"] + + def test_noop_when_no_mmproj(self): + cmd = ["llama-server", "-m", "/p/model.gguf", "-c", "4096", "--jinja"] + assert _strip(cmd) == cmd + + def test_returns_new_list(self): + cmd = ["llama-server", "--mmproj", "/p/mm.gguf"] + out = _strip(cmd) + assert out is not cmd + assert cmd[-1] == "/p/mm.gguf" # input untouched + + +class TestRetryContract: + """The two helpers compose into the load_model retry decision.""" + + def test_gemma4_failure_yields_valid_text_only_command(self): + # Old-llama.cpp projector abort -> retry, and the retry argv is a + # valid text-only launch (model + context kept, projector gone). + assert _detect(_GEMMA4_OLD_LLAMACPP_OUT) is True + retry_cmd = _strip(_VISION_CMD) + assert "--mmproj" not in retry_cmd + assert "-m" in retry_cmd and "--jinja" in retry_cmd + + def test_oom_does_not_retry_text_only(self): + # An OOM with --mmproj present must NOT be treated as a projector + # problem: load_model errors out instead of dropping vision. + assert _detect(_OOM_OUT) is False diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 4a8276adc0..b00cd7169e 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -77,9 +77,7 @@ def _enc_kv_string(key: str, value: str) -> bytes: def _enc_kv_uint32(key: str, value: int) -> bytes: - return ( - _enc_string(key) + struct.pack(" "off"). + # _already_in_target_state compares canonical requested modes; a vision + # backend with _requested_spec_mode="off" matches req "off" or None+vision. backend = _mtp_backend( _model_identifier = "unsloth/Qwen3-VL-4B-Instruct-GGUF", _is_vision = True, @@ -497,7 +492,7 @@ def test_unload_resets_nextn_predict_layers(): def _make_fake_llama_server(path: Path, help_text: str) -> Path: """Bash stub that prints `help_text` on --help.""" - path.write_text("#!/usr/bin/env bash\n" f"cat <<'EOF'\n{help_text}\nEOF\n") + path.write_text(f"#!/usr/bin/env bash\ncat <<'EOF'\n{help_text}\nEOF\n") path.chmod(0o755) return path @@ -532,8 +527,7 @@ def test_probe_server_capabilities_detects_renamed_mtp(tmp_path): # Renamed upstream: draft-mtp -> mtp. fake = _make_fake_llama_server( tmp_path / "llama-server", - "--spec-type [none|mtp|ngram-cache|ngram-simple|ngram-map-k|" - "ngram-map-k4v|ngram-mod]", + "--spec-type [none|mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod]", ) _clear_caps_cache() caps = LlamaCppBackend.probe_server_capabilities(str(fake)) @@ -619,9 +613,8 @@ def test_probe_detects_legacy_ngram_mod_flavor(tmp_path): @_NEEDS_BASH def test_probe_ignores_removal_stub_descriptions(tmp_path): - # Post-rename binary: legacy flags are present but with - # "argument has been removed" descriptions; must not be detected - # as legacy. + # Post-rename binary: legacy flags present but with "argument has been + # removed" descriptions; must not be detected as legacy. fake = _make_fake_llama_server(tmp_path / "llama-server", _POST_RENAME_HELP) _clear_caps_cache() caps = LlamaCppBackend.probe_server_capabilities(str(fake)) @@ -655,14 +648,7 @@ def test_build_ngram_mod_flags_new(): def test_build_ngram_mod_flags_legacy(): flags = _build_ngram_mod_flags({"ngram_mod_flavor": "legacy"}) - assert flags == [ - "--spec-ngram-size-n", - "24", - "--draft-min", - "48", - "--draft-max", - "64", - ] + assert flags == ["--spec-ngram-size-n", "24", "--draft-min", "48", "--draft-max", "64"] def test_build_ngram_mod_flags_empty_when_unsupported(): @@ -672,9 +658,7 @@ def test_build_ngram_mod_flags_empty_when_unsupported(): def test_build_ngram_mod_flags_respects_custom_values(): - flags = _build_ngram_mod_flags( - {"ngram_mod_flavor": "new"}, n_match = 16, n_min = 24, n_max = 32 - ) + flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"}, n_match = 16, n_min = 24, n_max = 32) assert flags == [ "--spec-ngram-mod-n-match", "16", @@ -796,16 +780,15 @@ def test_already_in_target_state_draft_n_max_ignored_when_not_mtp(): ) -# Sub-3B MTP gate -- tiny dense models regress with the MTP draft -# head, so load_model falls back to ngram-mod (when the binary supports -# it) instead of draft-mtp. The reload-skip mirror must follow the -# same fallback so a sub-3B reload-with-default does not bounce a -# correctly-configured ngram-mod / off backend. +# Sub-3B MTP gate -- tiny dense models regress with the MTP draft head, so +# load_model falls back to ngram-mod (when the binary supports it) instead of +# draft-mtp. The reload-skip mirror must follow the same fallback so a sub-3B +# reload-with-default doesn't bounce a correctly-configured ngram-mod/off backend. def _patch_probe(monkeypatch, ngram_supported): - """Force probe_server_capabilities to a deterministic result so - tests don't depend on whatever llama-server happens to be on PATH.""" + """Force probe_server_capabilities to a deterministic result so tests + don't depend on whatever llama-server is on PATH.""" fake = { "found": True, "mtp_token": "draft-mtp", @@ -826,11 +809,9 @@ def _patch_probe(monkeypatch, ngram_supported): ) -def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported( - monkeypatch, -): - # 0.8B MTP request -- load_model would have promoted to ngram-mod - # (no MTP head); reload check must match a ngram-mod backend. +def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported(monkeypatch): + # 0.8B MTP request -- load_model would have promoted to ngram-mod (no MTP + # head); reload check must match a ngram-mod backend. _patch_probe(monkeypatch, ngram_supported = True) backend = _mtp_backend( _model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF", @@ -902,8 +883,8 @@ def test_already_in_target_state_4b_mtp_request_promotes_as_before(monkeypatch): def test_already_in_target_state_2b_falls_back_to_ngram_below_threshold(monkeypatch): - # 2.0B is below the 3B threshold -> ngram-mod fallback, not - # draft-mtp. Clean-bench shows 2B regresses with draft-mtp. + # 2.0B is below the 3B threshold -> ngram-mod fallback, not draft-mtp. + # Clean-bench shows 2B regresses with draft-mtp. _patch_probe(monkeypatch, ngram_supported = True) backend = _mtp_backend( _model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF", @@ -1011,7 +992,12 @@ def test_canonicalize_spec_mode(value, expected): # ── _build_speculative_flags resolver matrix ────────────────────── -def _resolver_backend(monkeypatch, *, ngram_supported = True, mtp_token = "draft-mtp"): +def _resolver_backend( + monkeypatch, + *, + ngram_supported = True, + mtp_token = "draft-mtp", +): """Backend with a deterministic probe so the resolver is hermetic.""" fake = { "found": True, @@ -1032,9 +1018,9 @@ def _resolver_backend(monkeypatch, *, ngram_supported = True, mtp_token = "draft def _flags_dict(flags): - """Parse the spec-flag list into a small {flag: value} dict; collapses - repeated flags by keeping the last (only --spec-type can repeat and - never does in our resolver).""" + """Parse the spec-flag list into a {flag: value} dict; collapses repeated + flags by keeping the last (only --spec-type can repeat, and never does + in our resolver).""" out = {} i = 0 while i < len(flags): @@ -1070,8 +1056,8 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF" ("mtp", False, _MTP_MODEL, "draft-mtp", "3", False), # ── mtp forced on sub-3B: engage anyway ── ("mtp", True, _SUB_3B_MTP_MODEL, "draft-mtp", "2", False), - # ── mtp forced on non-MTP: engage anyway ── - ("mtp", True, _NON_MTP_MODEL, "draft-mtp", "2", False), + # ── mtp forced on non-MTP: default back (no head/drafter) ── + ("mtp", True, _NON_MTP_MODEL, None, None, False), # ── ngram forced: ngram-mod alone on BOTH platforms ── ("ngram", True, _MTP_MODEL, "ngram-mod", None, True), ("ngram", False, _MTP_MODEL, "ngram-mod", None, True), @@ -1080,6 +1066,8 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF" ("mtp+ngram", True, _MTP_MODEL, "ngram-mod,draft-mtp", "2", True), ("mtp+ngram", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True), ("mtp+ngram", True, _SUB_3B_MTP_MODEL, "ngram-mod,draft-mtp", "2", True), + # ── mtp+ngram forced on non-MTP: keep ngram, drop draft-mtp ── + ("mtp+ngram", True, _NON_MTP_MODEL, "ngram-mod", None, True), # ── off: nothing emitted ── ("off", True, _MTP_MODEL, None, None, False), ("off", False, _MTP_MODEL, None, None, False), @@ -1093,13 +1081,7 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF" ], ) def test_build_speculative_flags_matrix( - monkeypatch, - requested, - gpus, - model, - expect_spec_type, - expect_n_max, - expect_ngram_knobs, + monkeypatch, requested, gpus, model, expect_spec_type, expect_n_max, expect_ngram_knobs ): backend = _resolver_backend(monkeypatch) flags = backend._build_speculative_flags( @@ -1140,8 +1122,8 @@ def test_build_speculative_flags_user_extra_args_owns_spec_type(monkeypatch): gpus = True, binary = "/fake/llama-server", ) - # No flags emitted by the resolver -- the user's extra_args carries - # the --spec-type, and the resolver records requested_spec_mode = None. + # Resolver emits nothing -- the user's extra_args carries the --spec-type, + # and the resolver records requested_spec_mode = None. assert flags == [] assert backend.requested_spec_mode is None assert backend.speculative_type is None @@ -1194,7 +1176,285 @@ def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch): binary = "/fake/llama-server", ) assert "--spec-type" not in flags - # _speculative_type stays None (resolved emission was none), but - # _requested_spec_mode still reflects the user's choice. + # _speculative_type stays None (resolved emission was none); the user's + # choice is still reflected in _requested_spec_mode. assert backend.requested_spec_mode == "mtp" assert backend.speculative_type is None + + +def test_forced_mtp_on_non_mtp_model_defaults_back(monkeypatch): + # Forcing MTP on a model with no head/drafter must NOT emit draft-mtp: + # llama-server aborts on it ("failed to measure MTP context memory") + # rather than no-op'ing. Default back to --spec-default instead. + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _NON_MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert "--spec-type" not in flags + assert "--spec-default" in flags + assert backend.speculative_type == "default" + assert backend.requested_spec_mode == "mtp" + + +def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch): + # mtp+ngram on a non-MTP model drops the doomed draft-mtp chain but keeps + # the ngram half, which needs no head. + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "mtp+ngram", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _NON_MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "ngram-mod" + assert backend.speculative_type == "ngram-mod" + assert backend.requested_spec_mode == "mtp+ngram" + + +# ── Full named-repo resolver matrix (the shipping Studio families) ───── +# +# Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and +# gemma-4 (regular + QAT) GGUF repo, including the giant MoEs that stay +# resolver-only (122B-A10B / 397B-A17B). Expectations are derived from the +# same signals load_model uses -- _extract_model_size_b (active>effective> +# total, so E2B->2, A3B->3, A10B->10, A17B->17), _is_mtp_model_name, and the +# separate-drafter flag -- so each row mirrors what the loader emits on a +# B200 (GPU default, n=2). gemma carries no -MTP marker; its MTP comes from +# the root mtp-*.gguf drafter, modelled here by passing mtp_draft_path. +# +# auto_spec: "draft-mtp" = head/drafter engaged (>=3B MTP, or any size with a +# separate drafter); "ngram-mod" = embedded sub-3B drop (zero-VRAM); None = +# non-MTP -> llama-server --spec-default. + +_GEMMA_DRAFTER = "/snap/mtp-gemma-4-it.gguf" # stand-in separate drafter + +_REAL_REPO_MATRIX = [ + # repo, drafter, auto_spec, auto_ngram_knobs + ("unsloth/Qwen3.5-0.8B-MTP-GGUF", None, "ngram-mod", True), + ("unsloth/Qwen3.5-2B-MTP-GGUF", None, "ngram-mod", True), + ("unsloth/Qwen3.5-4B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-9B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-27B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-35B-A3B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-122B-A10B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-397B-A17B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-0.8B-GGUF", None, None, False), + ("unsloth/Qwen3.5-2B-GGUF", None, None, False), + ("unsloth/Qwen3.5-4B-GGUF", None, None, False), + ("unsloth/Qwen3.5-9B-GGUF", None, None, False), + # E2B is 2B but ships a separate drafter -> exempt from the sub-3B drop. + ("unsloth/gemma-4-E2B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-12b-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-26B-A4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-31B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E2B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-12b-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-26B-A4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-31B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), +] + + +def _resolve_real(monkeypatch, repo, drafter, mode): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = mode, + spec_draft_n_max = None, + extra_args = None, + model_identifier = repo, + model_path = None, + gpus = True, # B200 default + binary = "/fake/llama-server", + mtp_draft_path = drafter, + ) + return backend, flags, _flags_dict(flags) + + +@pytest.mark.parametrize( + "repo, drafter, auto_spec, auto_ngram_knobs", + _REAL_REPO_MATRIX, + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_auto_routing(monkeypatch, repo, drafter, auto_spec, auto_ngram_knobs): + # Auto is the default mode the dropdown ships with. + backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "auto") + if auto_spec is None: + # Non-MTP: no draft-mtp, hand off to llama-server's own default. + assert "--spec-type" not in parsed + assert "--spec-default" in flags + assert backend.speculative_type == "default" + elif auto_spec == "draft-mtp": + assert parsed.get("--spec-type") == "draft-mtp" + assert parsed.get("--spec-draft-n-max") == "2" + assert backend.speculative_type == "draft-mtp" + # gemma ships a separate drafter; Qwen bakes the head into the GGUF. + assert ( + (parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed) + ) + else: # ngram-mod (sub-3B MTP drop) + assert parsed.get("--spec-type") == "ngram-mod" + assert "--model-draft" not in parsed # draft head dropped + assert backend.speculative_type == "ngram-mod" + if auto_ngram_knobs: + assert "--spec-ngram-mod-n-match" in parsed + assert backend.requested_spec_mode == "auto" + + +@pytest.mark.parametrize( + "repo, drafter", + [(r[0], r[1]) for r in _REAL_REPO_MATRIX], + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_off_emits_nothing(monkeypatch, repo, drafter): + # Off must suppress speculative decoding for every family. + backend, flags, _ = _resolve_real(monkeypatch, repo, drafter, "off") + assert flags == [] + assert backend.speculative_type is None + assert backend.requested_spec_mode == "off" + + +@pytest.mark.parametrize( + "repo, drafter", + [(r[0], r[1]) for r in _REAL_REPO_MATRIX], + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_forced_mtp_never_aborts(monkeypatch, repo, drafter): + # Forcing MTP on the dropdown: real MTP models (name marker or separate + # drafter) engage draft-mtp even below 3B; non-MTP models default back to + # --spec-default instead of emitting a draft-mtp llama-server will abort on. + backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "mtp") + is_real_mtp = _is_mtp_model_name(repo) or bool(drafter) + if is_real_mtp: + assert parsed.get("--spec-type") == "draft-mtp" + assert backend.speculative_type == "draft-mtp" + assert ( + (parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed) + ) + else: + assert "--spec-type" not in parsed + assert "--spec-default" in flags + assert backend.speculative_type == "default" + assert backend.requested_spec_mode == "mtp" + + +# ── Sub-3B separate-drafter exemption (Gemma) ───────────────────────── +# +# The sub-3B MTP drop is an embedded-head cost (Qwen). A separate drafter +# (Gemma's root mtp-*.gguf) is a cheap standalone model that wins below 3B +# (B200 Q4_K_XL: gemma-4-E2B draft-mtp n=2 = 1.21x vs OFF), so it is exempt. + + +def test_sub3b_gemma_separate_drafter_engages_mtp(monkeypatch): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/gemma-4-E2B-it-GGUF", # 2B + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = "/snap/mtp-gemma-4-E2B-it.gguf", # separate drafter + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "draft-mtp" + assert parsed.get("--model-draft") == "/snap/mtp-gemma-4-E2B-it.gguf" + assert "--spec-ngram-mod-n-match" not in parsed + assert backend.speculative_type == "draft-mtp" + + +def test_sub3b_qwen_embedded_head_still_drops_to_ngram(monkeypatch): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF", # 2B, embedded head + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = None, # no separate drafter + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "ngram-mod" + assert "--model-draft" not in parsed + assert backend.speculative_type == "ngram-mod" + + +def test_auto_mode_drops_mtp_exempts_separate_drafter(): + from core.inference.llama_cpp import _auto_mode_drops_mtp + + assert _auto_mode_drops_mtp("auto", 2.0) is True + assert _auto_mode_drops_mtp("auto", 2.0, has_separate_drafter = True) is False + assert _auto_mode_drops_mtp("auto", 4.0) is False + assert _auto_mode_drops_mtp("mtp", 2.0) is False # forced engages regardless + + +# ── spec_fallback_reason (drives the "update llama.cpp" UI hint) ─────── + + +def test_spec_fallback_reason_set_when_binary_lacks_mtp(monkeypatch): + # Outdated llama-server with no mtp token: a forced MTP request can't emit + # draft-mtp, so record the reason for the UI update affordance. + backend = _resolver_backend(monkeypatch, mtp_token = None) + backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason == "binary_no_mtp" + + +def test_spec_fallback_reason_none_when_mtp_engages(monkeypatch): + backend = _resolver_backend(monkeypatch) + backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.speculative_type == "draft-mtp" + assert backend.spec_fallback_reason is None + + +def test_spec_fallback_reason_reset_on_off(monkeypatch): + # A subsequent off load must clear a stale reason. + backend = _resolver_backend(monkeypatch, mtp_token = None) + backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason == "binary_no_mtp" + backend._build_speculative_flags( + speculative_type = "off", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason is None diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py index b9f25faf88..10b1dc7ff6 100644 --- a/studio/backend/tests/test_llama_cpp_no_context_shift.py +++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py @@ -3,17 +3,15 @@ """``--no-context-shift`` launch-flag contract. -When llama-server runs with its default context-shift behavior, the UI -has no way to tell the user that the KV cache has been rotated -- -earlier turns silently vanish from the conversation. The Studio -backend always passes ``--no-context-shift`` so the server returns a +With llama-server's default context-shift behavior, the UI cannot tell the user +the KV cache was rotated -- earlier turns silently vanish from the conversation. +The Studio backend always passes ``--no-context-shift`` so the server returns a clean error instead, and the chat adapter can point the user at the ``Context Length`` input in the settings panel. -This file is a static read of the launch command: we ask -``LlamaCppBackend`` to assemble its ``cmd`` list and assert the flag -is always present. Testing via the real subprocess would require an -actual GGUF on disk, which is out of scope for the fast test suite. +This file statically reads the launch command: we ask ``LlamaCppBackend`` to +assemble its ``cmd`` list and assert the flag is present. Testing via the real +subprocess would need an actual GGUF on disk, out of scope for the fast suite. """ from __future__ import annotations @@ -68,11 +66,10 @@ from core.inference import llama_cpp as llama_cpp_module def _load_model_source() -> str: """Return the source of ``LlamaCppBackend.load_model``. - Using ``inspect.getsource`` instead of reading the file directly - scopes the assertions to the function that actually launches - llama-server, so neither the presence check nor the location check - can be fooled by a stray occurrence of ``"--no-context-shift"`` - elsewhere in the module. + Using ``inspect.getsource`` instead of reading the file scopes the assertions + to the function that launches llama-server, so neither the presence nor the + location check can be fooled by a stray ``"--no-context-shift"`` elsewhere in + the module. """ return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) @@ -80,10 +77,9 @@ def _load_model_source() -> str: def test_no_context_shift_is_in_load_model(): """The flag is part of the static launch-command template. - We check the source of ``load_model`` rather than mocking the whole - call chain (GPU probing, GGUF stat, etc.): the flag is written as - a literal in one place and any regression has to delete it, which - a text search will catch. + We check the source of ``load_model`` rather than mocking the whole call + chain (GPU probing, GGUF stat, etc.): the flag is a literal in one place and + any regression must delete it, which a text search catches. """ assert '"--no-context-shift"' in _load_model_source(), ( "llama-server must be launched with --no-context-shift so the " @@ -93,21 +89,19 @@ def test_no_context_shift_is_in_load_model(): def test_flag_sits_inside_the_base_cmd_list(): - """Pin the flag's location so a future refactor can't accidentally - move it into a branch that only fires on some code paths. + """Pin the flag's location so a refactor can't move it into a branch that + only fires on some code paths. - We slice from ``cmd = [`` to the first ``]`` at the same indent. - Using ``inspect.getsource`` means the function lives in its own - string and there are no siblings to worry about, so a plain - bracket search would also work -- anchoring on the trailing indent - just keeps the slice from wandering into a later expression if the - opening literal ever grows an in-line comment trailing it. + We slice from ``cmd = [`` to the first ``]`` at the same indent. Since + ``inspect.getsource`` gives the function its own string with no siblings, a + plain bracket search would also work -- anchoring on the trailing indent just + keeps the slice from wandering into a later expression if the opening literal + ever grows a trailing in-line comment. """ source = _load_model_source() start = source.find("cmd = [") assert start >= 0, "could not find the base cmd = [...] block" # Find the first line containing only ``]`` (possibly indented). - # Works for any indentation style the formatter picks. rest = source[start:] end_rel = -1 for line_start, line in _iter_lines_with_offset(rest): @@ -124,7 +118,7 @@ def test_flag_sits_inside_the_base_cmd_list(): "conditional branch -- otherwise some code paths would still " "run with silent context shift enabled." ) - # Also pin that it is next to -c / --ctx so the grouping makes sense. + # Pin that it sits next to -c / --ctx so the grouping makes sense. assert '"-c"' in block assert '"--flash-attn"' in block diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py new file mode 100644 index 0000000000..d87c05f2c6 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -0,0 +1,254 @@ +# 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 the post-launch /props context readback. + +llama-server's memory-fit step or --parallel slot split can allocate less +context than the requested -c while Studio keeps advertising the requested +value; clients sized to it then die on exceed_context_size_error 400s. +``_reconcile_effective_ctx_with_server`` must adopt the server's real +``default_generation_settings.n_ctx`` whenever it is smaller. + +Stubbed httpx; no subprocess, GPU, or network. Cross-platform. +""" + +from __future__ import annotations + +import json +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy/unavailable deps before importing the module under test. +# Mirrors test_llama_cpp_context_fit.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Prefer the real modules so importing this file first cannot poison later +# test modules with stubs; only stub what the environment genuinely lacks. +try: + import loggers # noqa: F401 +except ImportError: + _loggers_stub = _types.ModuleType("loggers") + _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) + sys.modules.setdefault("loggers", _loggers_stub) + +try: + import structlog # noqa: F401 +except ImportError: + sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "WriteError", + "HTTPError", + ): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + _httpx_stub.Timeout = _FakeTimeout + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + _httpx_stub.get = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError("unstubbed httpx.get")) + sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend +import core.inference.llama_cpp as llama_cpp_mod + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _FakeResponse: + def __init__( + self, + status_code = 200, + body = None, + ): + self.status_code = status_code + self._body = body or {} + + def json(self): + return self._body + + +def _make_backend(effective_ctx = 98304, port = 51234): + inst = LlamaCppBackend.__new__(LlamaCppBackend) + inst._port = port + inst._effective_context_length = effective_ctx + inst._context_length = 262144 + return inst + + +def _stub_props( + monkeypatch, + status_code = 200, + body = None, + exc = None, +): + def fake_get(url, timeout = None): + assert url.endswith("/props") + if exc is not None: + raise exc + return _FakeResponse(status_code, body) + + monkeypatch.setattr(llama_cpp_mod.httpx, "get", fake_get, raising = False) + + +# --------------------------------------------------------------------------- +# _query_server_n_ctx parsing +# --------------------------------------------------------------------------- + + +def test_query_n_ctx_reads_default_generation_settings(monkeypatch): + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 67584}}, + ) + assert _make_backend()._query_server_n_ctx() == 67584 + + +def test_query_n_ctx_non_200_returns_none(monkeypatch): + _stub_props(monkeypatch, status_code = 503) + assert _make_backend()._query_server_n_ctx() is None + + +def test_query_n_ctx_missing_key_returns_none(monkeypatch): + _stub_props(monkeypatch, body = {"default_generation_settings": {}}) + assert _make_backend()._query_server_n_ctx() is None + + +def test_query_n_ctx_swallows_transport_errors(monkeypatch): + _stub_props(monkeypatch, exc = RuntimeError("connection refused")) + assert _make_backend()._query_server_n_ctx() is None + + +# --------------------------------------------------------------------------- +# _reconcile_effective_ctx_with_server decisions +# --------------------------------------------------------------------------- + + +def test_fit_shrunk_ctx_overwrites_advertised_value(monkeypatch): + """The Nick repro: requested/advertised 98304, server really at 67584.""" + inst = _make_backend(effective_ctx = 98304) + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 67584}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 67584 + assert inst.context_length == 67584 + + +def test_matching_ctx_is_left_alone(monkeypatch): + inst = _make_backend(effective_ctx = 98304) + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 98304}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 98304 + + +def test_larger_server_ctx_does_not_inflate_advertised_value(monkeypatch): + """Never advertise more than the user asked for, even if the server could.""" + inst = _make_backend(effective_ctx = 32768) + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 65536}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 32768 + + +def test_unset_effective_ctx_adopts_server_value(monkeypatch): + inst = _make_backend(effective_ctx = None) + inst._context_length = None + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 40960}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 40960 + + +def test_props_failure_keeps_studio_value(monkeypatch): + """A flaky /props must never wipe the computed context.""" + inst = _make_backend(effective_ctx = 98304) + _stub_props(monkeypatch, exc = RuntimeError("boom")) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 98304 + + +# --------------------------------------------------------------------------- +# _ctx_integrity_flags: keep the per-request window equal to the advertised ctx +# --------------------------------------------------------------------------- + +_CAPS_ALL = {"supports_kv_unified": True, "supports_fit_ctx": True} +_CAPS_NONE = {"supports_kv_unified": False, "supports_fit_ctx": False} + + +def test_kv_unified_added_for_multi_slot(): + """Explicit --parallel N disables llama-server's auto-slots kv-unified + default, splitting -c into per-slot windows of -c/N; Studio must restore + the shared pool so one request can use the full advertised context.""" + flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL) + assert "--kv-unified" in flags + + +def test_kv_unified_skipped_for_single_slot_or_old_build(): + assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( + 1, False, 98304, 98304, _CAPS_ALL + ) + assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( + 4, False, 98304, 98304, _CAPS_NONE + ) + + +def test_fit_ctx_floors_explicit_request_under_fit(): + flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL) + assert flags[flags.index("--fit-ctx") + 1] == "98304" + + +def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support(): + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( + 1, False, 98304, 98304, _CAPS_ALL + ) + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL) + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, 98304, 98304, _CAPS_NONE + ) + + +def test_probe_missing_binary_reports_new_capabilities_false(): + info = LlamaCppBackend.probe_server_capabilities(binary = "/nonexistent/llama-server") + assert info["found"] is False + assert info["supports_kv_unified"] is False + assert info["supports_fit_ctx"] is False diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py new file mode 100644 index 0000000000..8d88a61ae3 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py @@ -0,0 +1,141 @@ +# 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 LlamaCppBackend._classify_llama_start_failure. + +When llama-server exits before becoming healthy, load_model turns its +captured stdout/stderr into a user-facing reason. A diffusion/image GGUF +(FLUX, Qwen-Image, ...) is a valid file with plenty of memory, so the +generic "invalid file or out of memory" message is misleading (issue +#5842). These tests pin the classification. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Match sibling tests' stubbing so the module imports in a lightweight +# env without fastapi. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +# Give the structlog stub a real get_logger: a bare ModuleType poisons +# sys.modules for later tests that call structlog.get_logger at import time. +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("structlog") +sys.modules.setdefault("structlog", _structlog_stub) +if not hasattr(sys.modules["structlog"], "get_logger"): + sys.modules["structlog"].get_logger = _structlog_stub.get_logger + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + +_classify = LlamaCppBackend._classify_llama_start_failure + +# Real llama-server failure lines (lower-cased downstream anyway). +_QWEN_IMAGE_OUT = ( + "load_model: loading model 'qwen-image-edit-2511-Q4_K_M.gguf'\n" + "llama_model_load: error loading model: unknown model architecture: 'qwen_image'\n" + "llama_model_load_from_file_impl: failed to load model" +) +_OOM_OUT = ( + "ggml_backend_cuda_buffer_type_alloc_buffer: allocating 12000.00 MiB on " + "device 0: cudaMalloc failed: out of memory" +) + + +class TestDiffusionArchitectures: + def test_qwen_image_routes_to_images_page(self): + msg = _classify(_QWEN_IMAGE_OUT, "/models/qwen-image.gguf", "local/qwen-image") + assert "diffusion" in msg.lower() + assert "Images page" in msg + assert "qwen_image" in msg + # Must NOT keep blaming memory / file validity. + assert "out of memory" not in msg.lower() + assert "enough memory" not in msg.lower() + + # Parametrize over the production set so new arches are auto-covered. + @pytest.mark.parametrize("arch", sorted(LlamaCppBackend._DIFFUSION_ARCHES)) + def test_every_diffusion_arch_is_recognised(self, arch): + out = f"error loading model: unknown model architecture: '{arch}'" + msg = _classify(out, f"/models/{arch}.gguf", f"local/{arch}") + assert "diffusion" in msg.lower() + assert "Images page" in msg + assert arch in msg + + +class TestUnsupportedNonDiffusionArchitecture: + def test_unknown_llm_arch_says_unsupported_not_oom(self): + out = "error loading model: unknown model architecture: 'some_new_llm'" + msg = _classify(out, "/models/x.gguf", "local/x") + assert "some_new_llm" in msg + assert "architecture" in msg.lower() + # Specific, not the misleading memory message. + assert "enough memory" not in msg.lower() + assert "diffusion" not in msg.lower() + + # Exact match: a chat arch merely containing a diffusion token (wan, + # sd1, flux, ...) must not be routed to the Images page. + @pytest.mark.parametrize( + "arch", + [ + "taiwan", # contains "wan" + "swan_llm", # contains "wan" + "fluxion", # contains "flux" + "sd1234", # contains "sd1" + "sd3_chat", # contains "sd3" + "aura2_text", # contains "aura" + "cosmos_reason", # contains "cosmos" + "qwen_image_text", # contains "qwen_image" + ], + ) + def test_arch_containing_diffusion_token_is_not_misrouted(self, arch): + out = f"error loading model: unknown model architecture: '{arch}'" + msg = _classify(out, f"/models/{arch}.gguf", f"local/{arch}") + assert arch in msg + assert "does not support" in msg.lower() + assert "diffusion" not in msg.lower() + assert "Images page" not in msg + + +class TestOllamaAndFallback: + _OLLAMA_GGUF = ( + f"/home/u/.ollama{__import__('os').sep}ollama_links" f"{__import__('os').sep}m.gguf" + ) + + def test_ollama_compat_message_still_works(self): + out = "llama_model_load: error loading model: key not found" + msg = _classify(out, self._OLLAMA_GGUF, "ollama/llama3") + assert "Ollama" in msg + + def test_ollama_unknown_arch_keeps_ollama_guidance(self): + # Ollama + non-diffusion unknown arch keeps the Ollama hint, not the + # generic llama.cpp "unsupported" message. + out = "error loading model: unknown model architecture: 'some_new_llm'" + msg = _classify(out, self._OLLAMA_GGUF, "ollama/some-new") + assert "Ollama" in msg + assert "directly through Ollama" in msg + assert "does not support" not in msg.lower() + + def test_ollama_diffusion_arch_still_routes_to_images(self): + # Diffusion routing wins over the Ollama hint. + out = "error loading model: unknown model architecture: 'flux'" + msg = _classify(out, self._OLLAMA_GGUF, "ollama/flux") + assert "diffusion" in msg.lower() + assert "Images page" in msg + + def test_generic_oom_keeps_memory_message(self): + msg = _classify(_OOM_OUT, "/models/big.gguf", "local/big") + assert "enough memory" in msg.lower() + assert "diffusion" not in msg.lower() + + def test_empty_output_is_safe(self): + msg = _classify("", None, None) + assert "llama-server failed to start" in msg diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py new file mode 100644 index 0000000000..fa583ef53d --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -0,0 +1,1151 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Focused tests for the GGUF llama.cpp agentic tool loop. + +These tests drive ``LlamaCppBackend.generate_chat_completion_with_tools`` +with fake llama-server SSE streams. They require no model, subprocess, GPU, +or network access. +""" + +from __future__ import annotations + +import contextlib +import copy +import json +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.llama_cpp import LlamaCppBackend + + +def _sse(delta: dict) -> str: + return "data: " + json.dumps({"choices": [{"index": 0, "delta": delta}]}) + "\n" + + +def _done() -> str: + return "data: [DONE]\n" + + +def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._process = object() + backend._healthy = True + backend._port = 48847 + backend._api_key = None + backend._effective_context_length = 4096 + backend._supports_reasoning = False + backend._reasoning_always_on = False + backend._reasoning_style = "enable_thinking" + backend._supports_preserve_thinking = False + + @contextlib.contextmanager + def fake_stream_with_retry( + _client, + _url, + payload, + _cancel_event, + headers = None, + ): + payloads.append(copy.deepcopy(payload)) + yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() + + def fake_iter_text_cancellable(response, _cancel_event): + yield from response.chunks + + monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) + monkeypatch.setattr(backend, "_iter_text_cancellable", fake_iter_text_cancellable) + return backend + + +def _tool_names(payload: dict) -> list[str]: + return [ + (tool.get("function") or {}).get("name") + for tool in payload.get("tools", []) + if (tool.get("function") or {}).get("name") + ] + + +def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): + """llama-server may emit content first and then native delta.tool_calls. + + Studio must not drop that tool call after it has streamed the preface. + """ + + tool_call_id = "call_render_late" + first_stream = [ + _sse({"content": "Here is the artifact.\n\n"}), + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": tool_call_id, + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "
red
", + "title": "Simple Red Square", + } + ), + }, + } + ] + } + ), + _done(), + ] + second_stream = [ + _sse({"content": "Done."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: Simple Red Square." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + ) + ) + + content_events = [e for e in events if e.get("type") == "content"] + assert content_events[0]["text"] == "Here is the artifact.\n\n" + + first_content_index = next( + i for i, event in enumerate(events) if event.get("type") == "content" + ) + actual_tool_start_index = next( + i + for i, event in enumerate(events) + if event.get("type") == "tool_start" and event.get("arguments", {}).get("code") + ) + assert first_content_index < actual_tool_start_index + + assert calls == [ + ( + "render_html", + { + "code": "
red
", + "title": "Simple Red Square", + }, + ) + ] + assert any(e.get("type") == "tool_end" and e.get("tool_name") == "render_html" for e in events) + + # The second llama-server request should include the assistant preface + # plus the structured tool call, preserving OpenAI-compatible ordering. + assert len(payloads) == 2 + assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"] + assert assistant_messages[-1]["content"] == "Here is the artifact.\n\n" + assert assistant_messages[-1]["tool_calls"][0]["id"] == tool_call_id + assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" + + +def test_repeat_render_html_nudge_is_not_user_visible_error(monkeypatch): + """A repeated render_html call is an internal no-op, not a visible card.""" + + first_stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_first", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "first", + "title": "First", + } + ), + }, + } + ] + } + ), + _done(), + ] + repeat_stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_repeat", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "repeat", + "title": "Repeat", + } + ), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Short note."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, repeat_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: First." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + }, + {"type": "function", "function": {"name": "web_search"}}, + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 2, + ) + ) + + assert calls == [ + ( + "render_html", + {"code": "first", "title": "First"}, + ) + ] + assert _tool_names(payloads[1]) == ["web_search"] + + actual_tool_starts = [ + event + for event in events + if event.get("type") == "tool_start" and event.get("arguments", {}).get("code") + ] + tool_ends = [ + event + for event in events + if event.get("type") == "tool_end" and event.get("tool_name") == "render_html" + ] + assert len(actual_tool_starts) == 1 + assert len(tool_ends) == 1 + + assert len(payloads) == 3 + render_tool_messages = [ + message + for message in payloads[2]["messages"] + if message.get("role") == "tool" and message.get("name") == "render_html" + ] + assert len(render_tool_messages) == 1 + internal_nudges = [ + message + for message in payloads[2]["messages"] + if message.get("role") == "user" + and "Do not call render_html again" in message.get("content", "") + ] + assert len(internal_nudges) == 1 + + +def test_render_html_success_drops_tool_schema_before_final_pass(monkeypatch): + first_stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_first", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps({"code": "ok"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Done."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + return "Rendered HTML artifact: Done." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Render this."}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + max_tool_iterations = 3, + ) + ) + + assert len(payloads) == 2 + assert "tools" not in payloads[1] + assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) + final_user_messages = [ + m.get("content", "") for m in payloads[1]["messages"] if m.get("role") == "user" + ] + assert not any("used all available tool calls" in message for message in final_user_messages) + + +def test_non_consecutive_duplicate_web_search_is_internal_noop(monkeypatch): + first_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_1", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + python_call = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_python", + "type": "function", + "function": { + "name": "python", + "arguments": json.dumps({"code": "print('ok')"}), + }, + } + ] + } + ), + _done(), + ] + duplicate_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_2", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer from gathered data."}), _done()] + payloads: list[dict] = [] + backend = _make_backend( + monkeypatch, + [first_search, python_call, duplicate_search, final_stream], + payloads, + ) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return f"ok:{name}" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search gpus in 2026 prices and use python"}], + tools = tools, + max_tool_iterations = 3, + ) + ) + + assert calls == [ + ("web_search", {"query": "gpu prices 2026"}), + ("python", {"code": "print('ok')"}), + ] + assert [ + event.get("tool_name") + for event in events + if event.get("type") == "tool_start" and event.get("tool_name") + ] == ["web_search", "python"] + assert [ + event.get("tool_name") + for event in events + if event.get("type") == "tool_end" and event.get("tool_name") + ] == ["web_search", "python"] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_search_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + assert len(payloads) == 4 + assert _tool_names(payloads[3]) == ["web_search", "python"] + duplicate_nudges = [ + message + for message in payloads[3]["messages"] + if message.get("role") == "user" + and "already completed successfully" in message.get("content", "") + ] + assert len(duplicate_nudges) == 1 + + +def test_duplicate_web_search_noop_allows_distinct_followup_tool(monkeypatch): + first_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_1", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + duplicate_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_2", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + python_call = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_python", + "type": "function", + "function": { + "name": "python", + "arguments": json.dumps({"code": "print('ok')"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer from gathered data."}), _done()] + payloads: list[dict] = [] + backend = _make_backend( + monkeypatch, + [first_search, duplicate_search, python_call, final_stream], + payloads, + ) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return f"ok:{name}" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search gpus in 2026 prices and use python"}], + tools = tools, + max_tool_iterations = 4, + ) + ) + + assert calls == [ + ("web_search", {"query": "gpu prices 2026"}), + ("python", {"code": "print('ok')"}), + ] + assert [ + event.get("tool_name") + for event in events + if event.get("type") == "tool_start" and event.get("tool_name") + ] == ["web_search", "python"] + assert [ + event.get("tool_name") + for event in events + if event.get("type") == "tool_end" and event.get("tool_name") + ] == ["web_search", "python"] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_search_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + assert len(payloads) == 4 + assert _tool_names(payloads[2]) == ["web_search", "python"] + duplicate_nudges = [ + message + for message in payloads[2]["messages"] + if message.get("role") == "user" + and "already completed successfully" in message.get("content", "") + ] + assert len(duplicate_nudges) == 1 + + +def test_repeated_duplicate_noop_transitions_to_final_pass(monkeypatch): + first_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_1", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + duplicate_one = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_2", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + duplicate_two = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_3", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer from first search."}), _done()] + payloads: list[dict] = [] + backend = _make_backend( + monkeypatch, + [first_search, duplicate_one, duplicate_two, final_stream], + payloads, + ) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search gpus"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 10, + ) + ) + + assert calls == [("web_search", {"query": "gpu prices 2026"})] + assert [event.get("tool_call_id") for event in events if event.get("type") == "tool_end"] == [ + "call_search_1" + ] + assert len(payloads) == 4 + assert "tools" not in payloads[-1] + assert any( + event.get("type") == "content" and event.get("text") == "Final answer from first search." + for event in events + ) + + +def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch): + same_turn_duplicates = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_1", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + }, + { + "index": 1, + "id": "call_search_2", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + }, + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [same_turn_duplicates, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "search-result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search gpus"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert calls == [("web_search", {"query": "gpu prices 2026"})] + assert [event.get("tool_call_id") for event in events if event.get("type") == "tool_end"] == [ + "call_search_1" + ] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_search_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + + +def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch): + same_turn_render_calls = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_html_1", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps({"code": "one"}), + }, + }, + { + "index": 1, + "id": "call_html_2", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps({"code": "two"}), + }, + }, + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [same_turn_render_calls, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: One." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "render html"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + max_tool_iterations = 2, + ) + ) + + assert calls == [("render_html", {"code": "one"})] + assert [ + event.get("tool_call_id") + for event in events + if event.get("type") == "tool_start" and not event.get("arguments") + ] == ["call_html_1"] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_html_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + assert len(payloads) == 2 + assert "tools" not in payloads[1] + render_nudges = [ + message + for message in payloads[1]["messages"] + if message.get("role") == "user" + and "Do not call render_html again" in message.get("content", "") + ] + assert len(render_nudges) == 1 + + +def test_disabled_tool_call_is_internal_noop(monkeypatch): + disabled_python = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_python_disabled", + "type": "function", + "function": { + "name": "python", + "arguments": json.dumps({"code": "print(1)"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "I cannot run Python here."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [disabled_python, final_stream], payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert not [event for event in events if event.get("type") in {"tool_start", "tool_end"}] + assert len(payloads) == 2 + disabled_nudges = [ + message + for message in payloads[1]["messages"] + if message.get("role") == "user" and "not enabled" in message.get("content", "") + ] + assert len(disabled_nudges) == 1 + + +def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch): + """After render_html succeeds, do not force another render_html call. + + The post-tool model pass can say it will use render_html again without + emitting a tool call. That should be accepted as a final model mistake, + not turned into repeated internal re-prompts after the artifact already + exists. + """ + + first_stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_first", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "first", + "title": "First", + } + ), + }, + } + ] + } + ), + _done(), + ] + post_tool_stream = [ + _sse({"content": "I will now use render_html again."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, post_tool_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: First." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + ) + ) + + assert len(payloads) == 2 + assert len(calls) == 1 + assert any( + event.get("type") == "content" and event.get("text") == "I will now use render_html again." + for event in events + ) + + +def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): + """No-tool re-prompt attempts should not concatenate into the UI.""" + + streams = [ + [_sse({"content": "I will use render_html now."}), _done()], + [_sse({"content": "Understood. I will use render_html now."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now."] + assert len(payloads) == 2 + + +def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch): + """A hidden forced re-prompt may fall back to a plain final answer.""" + + streams = [ + [_sse({"content": "I will use render_html now."}), _done()], + [ + _sse({"content": "No tool is needed. Final answer: use a red square."}), + _done(), + ], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ], + max_tool_iterations = 1, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == [ + "I will use render_html now.", + "No tool is needed. Final answer: use a red square.", + ] + assert len(payloads) == 2 + + +def test_internal_reprompt_disabled_when_auto_heal_disabled(monkeypatch): + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + auto_heal_tool_calls = False, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now."] + assert len(payloads) == 1 + + +def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatch): + streams = [ + [ + _sse( + { + "content": '{"name":"web_search","arguments":{"query":"x"}}' + } + ), + _done(), + ], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + auto_heal_tool_calls = False, + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "x"})] + assert not any( + event.get("type") == "content" and "" in event.get("text", "") + for event in events + ) + + +def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): + """Suppression ends once a forced re-prompt actually calls a tool.""" + + streams = [ + [_sse({"content": "I will use render_html now."}), _done()], + [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_forced", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "forced", + "title": "Forced", + } + ), + }, + } + ] + } + ), + _done(), + ], + [_sse({"content": "Final note after tool."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: Forced." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + ) + ) + + assert len(calls) == 1 + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now.", "Final note after tool."] + assert len(payloads) == 3 diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py new file mode 100644 index 0000000000..3b8f511a7c --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -0,0 +1,843 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic tests for the in-app llama.cpp update orchestration. + +No network, no real install: the GitHub release lookup and the installer +subprocess are both monkeypatched. Verifies detection (update_available) and +the apply flow (job lifecycle, installer invocation, post-swap re-read). +""" + +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path +from types import ModuleType + +import pytest + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import utils.llama_cpp_freshness as freshness # noqa: E402 +import utils.llama_cpp_update as upd # noqa: E402 + +MARKER = "UNSLOTH_PREBUILT_INFO.json" + + +class _FakeInstallerPopen: + """Stands in for the streamed installer process in _run_update.""" + + def __init__( + self, + cmd, + *, + returncode = 0, + lines = None, + on_start = None, + captured_kwargs = None, + **kwargs, + ): + if captured_kwargs is not None: + captured_kwargs.update(kwargs) + if on_start is not None: + on_start(list(cmd)) + self.returncode = returncode + self.stdout = iter(lines or []) + + def wait(self): + return self.returncode + + def kill(self): + pass + + +def _patch_installer_popen( + monkeypatch, + *, + returncode = 0, + lines = None, + on_start = None, + captured_kwargs = None, +): + monkeypatch.setattr( + upd.subprocess, + "Popen", + lambda cmd, **kw: _FakeInstallerPopen( + cmd, + returncode = returncode, + lines = lines, + on_start = on_start, + captured_kwargs = captured_kwargs, + **kw, + ), + ) + + +def _write_install( + dir_: Path, + tag: str, + repo: str = "unslothai/llama.cpp", + asset: str | None = None, + release_tag: str | None = None, +) -> str: + """Create a fake prebuilt install tree and return the llama-server path. + + ``asset`` is the bundle filename recorded in the marker; omit it to model an + older marker that predates asset-based ROCm forwarding (backward compat). + ``release_tag`` is the full release tag (e.g. a ``b9596-mix-`` mix + build); defaults to ``tag`` for a plain prebuilt.""" + bin_dir = dir_ / "build" / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + binary = bin_dir / "llama-server" + binary.write_text("#!/bin/sh\necho stub\n") + marker = { + "tag": tag, + "release_tag": release_tag or tag, + "published_repo": repo, + "installed_at_utc": "2020-01-01T00:00:00Z", + "bundle_profile": "cuda13-newer", + "runtime_line": "cuda13", + } + if asset is not None: + marker["asset"] = asset + (dir_ / MARKER).write_text(json.dumps(marker)) + return str(binary) + + +@pytest.fixture(autouse = True) +def _clean_state(monkeypatch, tmp_path): + freshness.reset_caches() + upd._reset_job_for_tests() + upd._resolve_memo.clear() + # Isolate the freshness disk cache so the suite never writes the real + # ~/.unsloth cache (the default when storage_roots can't be imported). + monkeypatch.setattr(freshness, "_cache_dir", lambda: tmp_path / ".freshness_cache") + # Deterministic markerless paths: no host-pinned binary, no custom dir. + monkeypatch.delenv("LLAMA_SERVER_PATH", raising = False) + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + # Never hit the network in these tests. + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + yield + freshness.reset_caches() + upd._reset_job_for_tests() + upd._resolve_memo.clear() + + +def _no_prebuilt(monkeypatch): + """Stub the host prebuilt probe to 'none available' (no source-build offer).""" + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: None) + + +def _prebuilt( + monkeypatch, + *, + repo = "unslothai/llama.cpp", + release_tag = "b9585", + llama_tag = None, + asset = None, +): + """Stub the host prebuilt probe to report an available prebuilt.""" + payload = { + "prebuilt_available": True, + "repo": repo, + "release_tag": release_tag, + "llama_tag": llama_tag or release_tag, + "asset": asset or f"llama-{release_tag}-bin-macos-arm64.tar.gz", + "install_kind": "macos-arm64", + } + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: payload) + + +def test_status_no_marker_no_prebuilt(monkeypatch, tmp_path): + # No marker AND no prebuilt available for the host -> unsupported (the genuine + # source-build-with-nothing-to-offer case). + binary = tmp_path / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker file alongside + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _no_prebuilt(monkeypatch) + st = upd.get_update_status() + assert st["supported"] is False + assert st["update_available"] is False + assert st["installed_tag"] is None + + +def test_status_source_build_offers_prebuilt(monkeypatch, tmp_path): + # Markerless source build with a prebuilt now available for the host: surface + # the update. Unknown installed version (source build) is treated as behind. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "b9585") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: None) + st = upd.get_update_status() + assert st["supported"] is True + assert st["update_available"] is True + assert st["source_build"] is True + assert st["latest_tag"] == "b9585" + assert st["published_repo"] == "unslothai/llama.cpp" + + +def test_status_source_build_compares_llama_tag(monkeypatch, tmp_path): + # release_tag may be a fork wrapper (v1.0); compare/display the upstream + # llama_tag (b9457) so a source build is not wrongly judged newer. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "v1.0", llama_tag = "b9457") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9000) + st = upd.get_update_status() + assert st["latest_tag"] == "b9457" # not the wrapper tag + assert st["update_available"] is True # 9000 < 9457 + + +def test_status_source_build_pinned_binary_not_offered(monkeypatch, tmp_path): + # LLAMA_SERVER_PATH pins a custom binary outside any llama.cpp dir; an apply + # could not take effect, so the button must not surface. + binary = tmp_path / "custom" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch) + st = upd.get_update_status() + assert st["supported"] is False + assert st["update_available"] is False + + +def test_llama_install_root_pinned_returns_none(monkeypatch, tmp_path): + binary = tmp_path / "custom" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + assert upd._llama_install_root(str(binary)) is None + + +def test_status_source_build_suppressed_when_newer(monkeypatch, tmp_path): + # A source build already newer than the latest prebuilt is not nagged. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "b9518") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600) + st = upd.get_update_status() + assert st["supported"] is True + assert st["update_available"] is False + assert st["installed_tag"] == "b9600" + + +def test_status_source_build_skips_probe_while_job_runs(monkeypatch, tmp_path): + # While the updater swaps the tree, status polls must not exec the binary + # being replaced (on Windows that exec can fail the installer's os.replace); + # the 3s poller only consumes job progress. + binary = tmp_path / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + probes = {"resolve": 0, "version": 0} + + def _count_resolve(*, force_refresh = False): + probes["resolve"] += 1 + return None + + def _count_version(b): + probes["version"] += 1 + return None + + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", _count_resolve) + monkeypatch.setattr(upd, "_installed_build_number", _count_version) + with upd._job_lock: + upd._job["state"] = upd._JOB_RUNNING + st = upd.get_update_status() + assert st["job"]["state"] == "running" + assert probes == {"resolve": 0, "version": 0} + + +def test_status_update_available(monkeypatch, tmp_path): + binary = _write_install(tmp_path, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + st = upd.get_update_status(force_refresh = True) + assert st["supported"] is True + assert st["installed_tag"] == "b9493" + assert st["latest_tag"] == "b9518" + assert st["update_available"] is True + + +def test_status_up_to_date(monkeypatch, tmp_path): + binary = _write_install(tmp_path, "b9518") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + st = upd.get_update_status(force_refresh = True) + assert st["installed_tag"] == "b9518" + assert st["latest_tag"] == "b9518" + assert st["update_available"] is False + + +def test_start_update_no_marker_no_prebuilt_refuses(monkeypatch, tmp_path): + binary = tmp_path / "llama-server" + binary.write_text("stub") # no marker + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _no_prebuilt(monkeypatch) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "no_prebuilt_available" + + +def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path): + # Markerless install + available prebuilt: install in place into the resolved + # root, with the asset-derived ROCm forwarding and the resolved repo. + install_dir = tmp_path / "llama.cpp" + binary = install_dir / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _prebuilt( + monkeypatch, repo = "unslothai/llama.cpp", asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz" + ) + + captured = {} + + class _Proc: + returncode = 0 + stdout = "installed" + stderr = "" + + def _fake_run(cmd, **kwargs): + cmd = list(cmd) + assert "--version" in cmd # only status polls still use run() + return _Proc() + + def _on_start(cmd): + captured["cmd"] = cmd + _write_install(install_dir, "b9585") # installer writes the marker + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + _patch_installer_popen(monkeypatch, on_start = _on_start) + + res = upd.start_update() + assert res["started"] is True, res + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + cmd = captured["cmd"] + assert "--install-dir" in cmd and str(install_dir) in cmd + assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd + assert "--llama-tag" in cmd and "latest" in cmd + assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x" + assert "--simple-policy" not in cmd and "--cpu-fallback" not in cmd + + +def test_start_update_happy_path(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + captured = {} + + class _Proc: + returncode = 0 + stdout = "installed" + stderr = "" + + def _on_start(cmd): + captured["cmd"] = cmd + # Simulate the installer writing a new marker with the latest tag. + _write_install(install_dir, "b9518") + + popen_kwargs: dict = {} + _patch_installer_popen( + monkeypatch, + lines = [ + "[llama-prebuilt] resolving release\n", + "Downloading llama.zip: 35.0% (12.0 MiB/35.0 MiB) at 9.0 MiB/s\n", + "Downloading llama.zip: 80.0% (28.0 MiB/35.0 MiB) at 9.0 MiB/s\n", + ], + on_start = _on_start, + captured_kwargs = popen_kwargs, + ) + + res = upd.start_update() + assert res["started"] is True + assert res["job"]["from_tag"] == "b9493" + assert res["job"]["progress"] == 0.0 + + # Wait for the background worker. + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert job["to_tag"] == "b9518" + # Installer was invoked with the resolved install dir + latest + repo. + assert "--install-dir" in captured["cmd"] + assert str(install_dir) in captured["cmd"] + assert "--llama-tag" in captured["cmd"] and "latest" in captured["cmd"] + assert "unslothai/llama.cpp" in captured["cmd"] + # Progress lines were parsed and success pins progress at 1.0. + assert job["progress"] == 1.0 + # The worker asks the installer for fine-grained progress milestones. + assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5" + + +def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + _patch_installer_popen(monkeypatch, returncode = 2, lines = ["boom: network error\n"]) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "error" + assert "boom" in (job["error"] or "") + + +# --- installer-argument construction (mirrors the post-#5963 setup scripts) --- + + +def test_rocm_install_args_lemonade_gfx(): + # Lemonade HIP app bundle: gfx family lives in the asset name. + assert upd._rocm_install_args("app-b9585-linux-x64-rocm-gfx110X.tar.gz") == [ + "--rocm-gfx", + "gfx110x", + ] + assert upd._rocm_install_args("app-b9585-windows-x64-rocm-gfx1150.zip") == [ + "--rocm-gfx", + "gfx1150", + ] + + +def test_rocm_install_args_fork_version_bundle(): + # Fork ROCm bundles encode a ROCm version, not a gfx -> forward --has-rocm. + assert upd._rocm_install_args("llama-b9334-bin-ubuntu-rocm-6.4-x64.tar.gz") == ["--has-rocm"] + + +def test_rocm_install_args_windows_hip(): + assert upd._rocm_install_args("llama-b9334-bin-win-hip-radeon-x64.zip") == ["--has-rocm"] + + +def test_rocm_install_args_non_rocm_and_missing(): + assert upd._rocm_install_args("llama-b9334-bin-ubuntu-x64.tar.gz") == [] + assert upd._rocm_install_args("app-b9585-linux-x64-cuda13.tar.gz") == [] + assert upd._rocm_install_args(None) == [] + + +def _capture_install_cmd( + monkeypatch, + tmp_path, + *, + tag = "b9493", + repo = "unslothai/llama.cpp", + asset = None, + latest = "b9518", +) -> list: + """Run start_update() with the installer subprocess stubbed; return the argv.""" + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, tag, repo = repo, asset = asset) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest) + + captured = {} + + class _Proc: + returncode = 0 + stdout = "installed" + stderr = "" + + def _fake_run(cmd, **kwargs): + cmd = list(cmd) + assert "--version" in cmd # only status polls still use run() + return _Proc() + + def _on_start(cmd): + captured["cmd"] = cmd + _write_install(install_dir, latest, repo = repo, asset = asset) + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + _patch_installer_popen(monkeypatch, on_start = _on_start) + + res = upd.start_update() + assert res["started"] is True, res + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + return captured.get("cmd", []) + + +def test_install_cmd_rocm_marker_forwards_gfx(monkeypatch, tmp_path): + cmd = _capture_install_cmd( + monkeypatch, tmp_path, asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz" + ) + assert "--rocm-gfx" in cmd + assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x" + assert "--has-rocm" not in cmd + assert "--cpu-fallback" not in cmd + assert "--simple-policy" not in cmd + assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd + + +def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): + cmd = _capture_install_cmd( + monkeypatch, tmp_path, asset = "llama-b9334-bin-ubuntu-rocm-6.4-x64.tar.gz" + ) + assert "--has-rocm" in cmd + assert "--rocm-gfx" not in cmd + + +def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): + # CPU installs come from ggml-org. Re-running into the same install-dir/repo + # reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU + # detection) is reserved for setup.sh's arm64 rescue and must not appear here. + cmd = _capture_install_cmd( + monkeypatch, + tmp_path, + repo = "ggml-org/llama.cpp", + asset = "llama-b9334-bin-ubuntu-x64.tar.gz", + ) + assert "--cpu-fallback" not in cmd + assert "--rocm-gfx" not in cmd + assert "--has-rocm" not in cmd + assert "--simple-policy" not in cmd + assert "--published-repo" in cmd and "ggml-org/llama.cpp" in cmd + + +def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tmp_path): + # Marker without an asset field (older install): no ROCm flags, no crash, and + # never the obsolete --simple-policy that #5963 removed from setup. + cmd = _capture_install_cmd(monkeypatch, tmp_path, asset = None) + assert "--simple-policy" not in cmd + assert "--rocm-gfx" not in cmd + assert "--has-rocm" not in cmd + assert "--cpu-fallback" not in cmd + + +# --- refusal + maintenance-state coordination --- + + +def test_start_update_already_running_refuses(monkeypatch, tmp_path): + binary = _write_install(tmp_path / "llama.cpp", "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + with upd._job_lock: + upd._job.update(state = upd._JOB_RUNNING) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "already_running" + + +def test_start_update_installer_missing_refuses(monkeypatch, tmp_path): + binary = _write_install(tmp_path / "llama.cpp", "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: None) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "installer_missing" + + +class _FakeBackend: + """Minimal stand-in for LlamaCppBackend's update-coordination surface.""" + + def __init__(self): + import threading + + self._serial_load_lock = threading.Lock() + self._llama_update_in_progress = False + self.is_active = True + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + +def _inject_backend(monkeypatch, backend): + routes_pkg = ModuleType("routes") + routes_pkg.__path__ = [] + inference_mod = ModuleType("routes.inference") + inference_mod.get_llama_cpp_backend = lambda: backend + monkeypatch.setitem(sys.modules, "routes", routes_pkg) + monkeypatch.setitem(sys.modules, "routes.inference", inference_mod) + + +def test_update_sets_maintenance_flag_and_unloads(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + backend = _FakeBackend() + _inject_backend(monkeypatch, backend) + + seen = {} + + def _on_start(cmd): + # The maintenance flag must be set while the installer runs. + seen["flag_during_install"] = backend._llama_update_in_progress + _write_install(install_dir, "b9518") + + _patch_installer_popen(monkeypatch, on_start = _on_start) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + + assert backend.unloaded is True + assert seen.get("flag_during_install") is True + # Cleared in the finally so model loads work again after the swap. + assert backend._llama_update_in_progress is False + + +def test_update_clears_maintenance_flag_on_installer_failure(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + backend = _FakeBackend() + _inject_backend(monkeypatch, backend) + + _patch_installer_popen(monkeypatch, returncode = 1, lines = ["boom\n"]) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + assert upd.get_update_status()["job"]["state"] == "error" + assert backend._llama_update_in_progress is False + + +def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + def _raise(): + raise RuntimeError("no backend") + + inference_mod = ModuleType("routes.inference") + inference_mod.get_llama_cpp_backend = lambda: _raise() + routes_pkg = ModuleType("routes") + routes_pkg.__path__ = [] + monkeypatch.setitem(sys.modules, "routes", routes_pkg) + monkeypatch.setitem(sys.modules, "routes.inference", inference_mod) + + _patch_installer_popen(monkeypatch, on_start = lambda cmd: _write_install(install_dir, "b9518")) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + + +# --- markerless helper units --- + + +def test_resolve_prebuilt_parses_and_caches(monkeypatch, tmp_path): + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + calls = {"n": 0} + + class _Proc: + returncode = 0 + # stderr noise plus the JSON line on stdout (installer logs to stderr). + stdout = ( + '{"prebuilt_available": true, "repo": "unslothai/llama.cpp", "release_tag": "b9585"}' + ) + stderr = "[llama-prebuilt] some log\n" + + def _fake_run(cmd, **kwargs): + calls["n"] += 1 + assert "--resolve-prebuilt" in cmd + return _Proc() + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + res = upd._resolve_prebuilt_for_host() + assert res["prebuilt_available"] is True and res["release_tag"] == "b9585" + # Second call is memoized (no second subprocess). + upd._resolve_prebuilt_for_host() + assert calls["n"] == 1 + + +def test_resolve_prebuilt_fails_open(monkeypatch, tmp_path): + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + + def _boom(cmd, **kwargs): + raise OSError("subprocess failed") + + monkeypatch.setattr(upd.subprocess, "run", _boom) + assert upd._resolve_prebuilt_for_host() is None + # Failures are not cached: a later success is observed. + + class _Proc: + returncode = 0 + stdout = '{"prebuilt_available": false}' + stderr = "" + + monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + assert upd._resolve_prebuilt_for_host() == {"prebuilt_available": False} + + +def test_installed_build_number(monkeypatch): + def _ver(text): + class _Proc: + returncode = 0 + stdout = "" + stderr = text + + monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + return upd._installed_build_number("/bin/llama-server") + + assert _ver("version: 9585 (abc1234)\nbuilt with clang\n") == 9585 + assert _ver("version: 1 (deadbee)\n") is None # source build without tags + assert _ver("no version here") is None + assert upd._installed_build_number(None) is None + + +def test_llama_install_root_finds_llama_cpp_ancestor(monkeypatch, tmp_path): + root = tmp_path / "llama.cpp" + binary = root / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) == root + + +def test_llama_install_root_unmanaged_path_returns_none(monkeypatch, tmp_path): + # A binary on PATH (no marker, no env pin, no llama.cpp ancestor) is foreign: + # installing elsewhere would not replace it, so report no manageable root. + binary = tmp_path / "usr" / "local" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) is None + + +def test_llama_install_root_unsloth_env_dir(monkeypatch, tmp_path): + # UNSLOTH_LLAMA_CPP_PATH dir holding the active binary is the managed root. + root = tmp_path / "vendor" / "llama" + binary = root / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("UNSLOTH_LLAMA_CPP_PATH", str(root)) + assert upd._llama_install_root(str(binary)) == root + + +def test_llama_install_root_ignores_inactive_env_root(monkeypatch, tmp_path): + # UNSLOTH_LLAMA_CPP_PATH set but the active binary is not under it: do not + # target the stale env root, resolve from the binary's own llama.cpp tree. + inactive = tmp_path / "custom-empty" + inactive.mkdir() + active = tmp_path / "llama.cpp" + binary = active / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("UNSLOTH_LLAMA_CPP_PATH", str(inactive)) + assert upd._llama_install_root(str(binary)) == active + + +def test_llama_install_root_refuses_pinned_checkout_under_llama_cpp(monkeypatch, tmp_path): + # The LLAMA_SERVER_PATH pin guard must run before the ancestor scan, or a + # user's own llama.cpp checkout could be handed to the installer. + root = tmp_path / "my-project" / "llama.cpp" + binary = root / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) is None + + +def test_start_update_source_build_refuses_when_newer(monkeypatch, tmp_path): + # A direct POST on a source build already newer than the prebuilt must not + # downgrade it; start_update mirrors the detection suppression. + install_dir = tmp_path / "llama.cpp" + binary = install_dir / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _prebuilt(monkeypatch, release_tag = "b9518") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "up_to_date" + + +# --- mix-tag detection + apply guard (the reported banner bug) --- + + +def test_status_not_offered_on_mix_latest(monkeypatch, tmp_path): + # Installed the mix latest; GitHub latest is that same full tag -> no banner. + binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453" + ) + st = upd.get_update_status() + assert st["update_available"] is False + assert st["installed_tag"] == "b9596" + assert st["latest_tag"] == "b9596-mix-e6f2453" + + +def test_status_not_offered_when_latest_lags(monkeypatch, tmp_path): + # A lagging latest (older build than installed) must never be offered. + binary = _write_install(tmp_path / "llama.cpp", "b9585") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + st = upd.get_update_status() + assert st["update_available"] is False + + +def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path): + # A direct POST / stale banner must not reinstall when already on the latest. + binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453" + ) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "up_to_date" diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py index bcf2eb1683..1ba6c9f7b5 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -3,10 +3,10 @@ """Tests for LlamaCppBackend._wait_for_health resilience. -The probe loop must swallow transient httpx errors and fall through to -the subprocess.poll() branch so a crashed llama-server surfaces a -structured "exited with code X" log instead of bubbling an opaque -exception up to the /api/inference/load route. +The probe loop must swallow transient httpx errors and fall through to the +subprocess.poll() branch so a crashed llama-server surfaces a structured +"exited with code X" log instead of bubbling an opaque exception up to the +/api/inference/load route. """ from __future__ import annotations @@ -22,8 +22,7 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Match the stubbing pattern in sibling tests so the module imports in -# a lightweight env without fastapi. +# Mirror sibling tests' stubbing so the module imports without fastapi. _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) @@ -33,11 +32,10 @@ import httpx # noqa: E402 from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 -# Sibling tests in this directory install lightweight httpx stubs via -# sys.modules.setdefault. When collected together, our `httpx` symbol -# may be one of those stubs, which lacks `get`. Ensure the production -# code finds a working `httpx.get` and the standard exception types -# regardless of collection order by adding the missing attributes. +# Sibling tests install lightweight httpx stubs via sys.modules.setdefault. +# When collected together, our `httpx` may be such a stub lacking `get`. Add +# the missing attributes so production code finds a working `httpx.get` and +# the standard exception types regardless of collection order. if not hasattr(httpx, "get"): httpx.get = None # placeholder; every test below monkeypatches it for _exc_name in ( @@ -52,9 +50,7 @@ for _exc_name in ( def _make_backend(port: int = 12345) -> LlamaCppBackend: - """Build a barebones LlamaCppBackend instance with only the - attributes _wait_for_health touches. Bypasses __init__ so we do not - pull in the full subprocess + logging stack.""" + """Barebones LlamaCppBackend with only the attributes _wait_for_health touches (bypasses __init__).""" b = LlamaCppBackend.__new__(LlamaCppBackend) b._port = port b._stdout_thread = None @@ -72,14 +68,9 @@ class TestWaitForHealthResilience: assert b._wait_for_health(timeout = 1.0, interval = 0.01) is True def test_read_error_loops_to_subprocess_poll(self, monkeypatch): - """WinError 10054 maps to httpx.ReadError. The loop must swallow - it and the next iteration must detect the dead subprocess via - poll() != None, returning False with a structured exit-code log - instead of bubbling the ReadError.""" + """WinError 10054 (httpx.ReadError) must be swallowed; the next iteration sees the dead subprocess and returns False with a structured exit-code log.""" b = _make_backend() - # First iteration: process alive (so we reach the httpx probe). - # Second iteration: process has exited (so we hit the structured - # exit-code branch and return False). + # Iter 1: alive (reach probe); iter 2: exited (exit-code branch -> False). b._process.poll.side_effect = [None, 1] b._process.returncode = 1 b._stdout_lines = ["llama-server: ggml-cuda.dll failed to load"] @@ -89,12 +80,12 @@ class TestWaitForHealthResilience: monkeypatch.setattr(httpx, "get", raise_read_error) assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False - # Both iterations of the loop ran -- the ReadError did not bubble. + # Both loop iterations ran -- the ReadError did not bubble. assert b._process.poll.call_count >= 2 def test_remote_protocol_error_also_swallowed(self, monkeypatch): - """Partial / malformed response on the probe (server crashed - mid-headers) raises RemoteProtocolError -- also non-fatal.""" + """A partial/malformed probe response (server crashed mid-headers) + raises RemoteProtocolError -- also non-fatal.""" b = _make_backend() b._process.poll.side_effect = [None, -1] b._process.returncode = -1 @@ -121,8 +112,8 @@ class TestWaitForHealthResilience: assert b._process.poll.call_count >= 2 def test_connect_error_swallowed_until_success(self, monkeypatch): - """Sanity: existing ConnectError swallowing still works -- the - loop retries until llama-server eventually answers 200.""" + """Sanity: existing ConnectError swallowing still works -- the loop + retries until llama-server answers 200.""" b = _make_backend() b._process.poll.return_value = None calls = {"n": 0} @@ -139,8 +130,8 @@ class TestWaitForHealthResilience: assert calls["n"] >= 3 def test_dead_process_before_probe_returns_false(self, monkeypatch): - """If poll() != None on entry, _wait_for_health must return - False immediately without calling httpx at all.""" + """poll() != None on entry: _wait_for_health returns False + immediately without calling httpx.""" b = _make_backend() b._process.poll.return_value = 137 b._process.returncode = 137 @@ -154,3 +145,111 @@ class TestWaitForHealthResilience: monkeypatch.setattr(httpx, "get", should_not_be_called) assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False assert called["n"] == 0 + + +class TestCrashLogTail: + """The "exited with code X" log must keep the TAIL of the output. + + Crash diagnostics (abort reason, ROCm/CUDA error text) print last, + after the long startup banner; head truncation has cut off exactly + the diagnostic line in field reports (gfx1151 fit-step abort).""" + + @staticmethod + def _capture_error_logs(monkeypatch) -> list: + """Capture module-logger .error() messages directly -- immune to + whatever logging/structlog config sibling test modules installed.""" + import core.inference.llama_cpp as _llama_mod + + records: list = [] + fake_logger = mock.Mock() + fake_logger.error = mock.Mock(side_effect = lambda msg, *a, **k: records.append(msg)) + monkeypatch.setattr(_llama_mod, "logger", fake_logger) + return records + + def test_crash_log_keeps_tail_not_head(self, monkeypatch): + records = self._capture_error_logs(monkeypatch) + b = _make_backend() + b._process.poll.return_value = 1 + b._process.returncode = 1 + # >2000 chars of banner, diagnostic on the final line. + banner = [f"load_model: tensor blk.{i} buffer ROCm0" for i in range(80)] + diagnostic = "ggml-cuda.cu:103: ROCm error: out of memory" + b._stdout_lines = banner + [diagnostic] + + assert b._wait_for_health(timeout = 1.0, interval = 0.01) is False + + crash_logs = [m for m in records if "exited with code" in m] + assert crash_logs, "crash must produce an exited-with-code log" + assert diagnostic in crash_logs[-1] + assert "Output (tail)" in crash_logs[-1] + # The head of the banner must be the part sacrificed to truncation. + assert "blk.0 buffer" not in crash_logs[-1] + + def test_crash_log_mentions_log_file_when_present(self, monkeypatch): + records = self._capture_error_logs(monkeypatch) + b = _make_backend() + b._process.poll.return_value = 1 + b._process.returncode = 1 + b._stdout_lines = ["boom"] + b._llama_log_path = Path("C:/logs/llama-123-port-1234.log") + + assert b._wait_for_health(timeout = 1.0, interval = 0.01) is False + + crash_logs = [m for m in records if "exited with code" in m] + assert crash_logs and "llama-123-port-1234.log" in crash_logs[-1] + + +class TestRetryLogFilenameUnique: + """The --fit off retry can respawn within the same epoch second; the log + filename must carry the attempt index or the second open ("w") truncates + the crash log the retry warning just referenced (found by simulation: + frozen time.time -> single file, crash evidence gone).""" + + def test_log_name_includes_attempt_index(self): + src = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + ).read_text(encoding = "utf-8") + assert "-try{_spawn_attempt}.log" in src + + +class TestFitOffRetryEligible: + """Gate for the one-shot --fit off startup-crash retry. + + Retry only when Studio's own VRAM math placed the model and nothing + on the command line chose the fit mode explicitly.""" + + def test_eligible_for_plain_ngl_launch(self): + cmd = ["llama-server", "-m", "x.gguf", "-ngl", "-1", "--jinja"] + assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is True + + def test_not_eligible_when_use_fit(self): + cmd = ["llama-server", "-m", "x.gguf", "--fit", "on"] + assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = True) is False + + @pytest.mark.parametrize( + "fit_args", + [ + ["--fit", "on"], + ["--fit", "off"], + ["-fit", "off"], + ["--fit=on"], + ["-fit=off"], + ], + ) + def test_not_eligible_with_explicit_fit_flag(self, fit_args): + cmd = ["llama-server", "-m", "x.gguf", *fit_args] + assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is False + + @pytest.mark.parametrize( + "tuning_args", + [ + ["--fit-ctx", "8192"], + ["--fit-target", "1024"], + ["-fitc", "4096"], + ["-fitt", "512"], + ["--fit-ctx=8192"], + ], + ) + def test_fit_tuning_flags_do_not_block_retry(self, tuning_args): + cmd = ["llama-server", "-m", "x.gguf", *tuning_args] + assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is True diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py index 00295d6283..493bb93e8c 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py @@ -3,9 +3,9 @@ """``_wait_for_vram_settle`` helper contract. -Pins the bounded poll over ``_get_gpu_free_memory`` that bridges the -kill -> spawn VRAM-reclaim window. Patches ``_get_gpu_free_memory``; -no real llama-server or nvidia-smi involved. +Pins the bounded poll over ``_get_gpu_free_memory`` bridging the kill -> spawn +VRAM-reclaim window. Patches ``_get_gpu_free_memory``; no real llama-server or +nvidia-smi involved. """ from __future__ import annotations @@ -19,10 +19,7 @@ from unittest.mock import patch import pytest -# --------------------------------------------------------------------------- -# Same external-dep stubs as the other llama_cpp tests so this module -# imports cleanly without httpx / structlog / loggers installed. -# --------------------------------------------------------------------------- +# External-dep stubs so this module imports without httpx / structlog / loggers. _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) @@ -34,8 +31,7 @@ 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) -# Ensure get_logger is set even if a previous test module already -# inserted a bare ``structlog`` stub via ``setdefault``. +# Set get_logger even if a prior test inserted a bare ``structlog`` stub. if not hasattr(sys.modules["structlog"], "get_logger"): sys.modules["structlog"].get_logger = _structlog_stub.get_logger @@ -74,8 +70,8 @@ def _patch_probe(samples): """Patch ``_get_gpu_free_memory`` to yield ``samples`` in order. Each entry is a list[(idx, free_mib)], a callable, or an exception - (instance or class). Calls past the end repeat the last entry so - tests can assert "stopped polling" via the call count. + (instance or class). Calls past the end repeat the last entry so tests + can assert "stopped polling" via the call count. """ state = {"i": 0, "calls": 0} @@ -112,8 +108,8 @@ def _kw(**extra): def test_cold_start_returns_immediately_without_probing(): - """Default ``since_kill=0.0`` is cold-start: no kill recorded, - helper short-circuits without ever invoking the probe.""" + """Default ``since_kill=0.0`` is cold-start: no kill recorded, so the + helper short-circuits without invoking the probe.""" ctx, state = _patch_probe([[(0, 10000)], [(0, 10000)]]) with ctx: start = time.monotonic() @@ -131,9 +127,7 @@ def test_stale_kill_skips_wait(): LlamaCppBackend._wait_for_vram_settle( **_kw(since_kill = long_ago, max_wait = 2.0, interval = 0.25) ) - assert ( - state["calls"] == 0 - ), "kill older than _VRAM_SETTLE_WINDOW_S must skip the wait" + assert state["calls"] == 0, "kill older than _VRAM_SETTLE_WINDOW_S must skip the wait" def test_empty_first_sample_returns_immediately(): @@ -156,8 +150,8 @@ def test_first_probe_raises_returns_without_polling(): def test_two_consecutive_samples_within_tolerance_settles(): - """The reclaim ramp from 10000 → 11500 → 11550: third sample within - 256 MiB of the second so the helper returns after exactly three probes.""" + """Reclaim ramp 10000 → 11500 → 11550: third sample within 256 MiB of + the second, so the helper returns after exactly three probes.""" ctx, state = _patch_probe( [ [(0, 10000)], @@ -170,7 +164,7 @@ def test_two_consecutive_samples_within_tolerance_settles(): LlamaCppBackend._wait_for_vram_settle(**_kw(max_wait = 2.0, interval = 0.05)) elapsed = time.monotonic() - start assert state["calls"] == 3 - # interval * 2 sleeps = 0.10; allow generous slack for scheduler jitter. + # interval * 2 sleeps = 0.10; allow slack for scheduler jitter. assert elapsed < 1.0 @@ -201,7 +195,7 @@ def test_max_wait_respected_when_never_settles(): start = time.monotonic() LlamaCppBackend._wait_for_vram_settle(**_kw(max_wait = 0.5, interval = 0.1)) elapsed = time.monotonic() - start - # We must stop near max_wait, not run forever. Generous upper bound for CI. + # Must stop near max_wait, not run forever. Generous upper bound for CI. assert 0.3 <= elapsed < 2.0, f"helper ignored max_wait: elapsed={elapsed:.3f}s" @@ -219,11 +213,9 @@ def test_max_wait_respected_when_probe_is_slow(): **_kw(max_wait = 0.4, interval = 0.25), ) elapsed = time.monotonic() - start - # First probe (0.30 s) + at most one short clipped sleep + bail. - # Hard cap well below the old behaviour of 0.30 + 0.25 + 0.30 = 0.85. - assert ( - elapsed < 0.85 - ), f"helper exceeded the deadline due to slow probes: {elapsed:.3f}s" + # First probe (0.30 s) + at most one clipped sleep + bail. + # Hard cap well below the old 0.30 + 0.25 + 0.30 = 0.85. + assert elapsed < 0.85, f"helper exceeded the deadline due to slow probes: {elapsed:.3f}s" def test_gpu_index_set_change_returns(): @@ -268,21 +260,19 @@ def test_tolerance_two_percent_for_large_cards(): def test_load_model_calls_helper_outside_lock_and_uses_last_kill_timestamp(): - """Pin the call site: outside Phase 3 lock, gated on the timestamp, - no ``had_live_process`` in-band flag regression. Mirrors the - ``inspect.getsource`` pattern from ``test_llama_cpp_no_context_shift``. - """ + """Pin the call site: outside Phase 3 lock, gated on the timestamp, no + ``had_live_process`` in-band flag regression.""" import inspect src = inspect.getsource(LlamaCppBackend.load_model) assert "_wait_for_vram_settle" in src assert "since_kill" in src assert "self._last_kill_monotonic" in src - # Must be invoked before Phase 3's broad lock so /unload, /cancel, - # /status are not blocked during the wait. + # Must run before Phase 3's broad lock so /unload, /cancel, /status + # are not blocked during the wait. assert src.index("_wait_for_vram_settle") < src.index("# ── Phase 3:") - # An in-band ``had_live_process`` flag would silently regress the - # frontend /unload+/load Apply path; use the timestamp instead. + # An in-band ``had_live_process`` flag would regress the frontend + # /unload+/load Apply path; use the timestamp instead. assert "had_live_process" not in src diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py index 7d4719c0e7..957de4bad6 100644 --- a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py +++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py @@ -4,9 +4,9 @@ """Tests for the Windows pip-nvidia DLL dir resolver. Studio installs torch with bundled CUDA wheels (nvidia-cuda-runtime-cu13, -nvidia-cublas-cu13, etc.) and the prebuilt llama-server.exe must find -those DLLs at runtime to load CUDA. Mirrors the Linux LD_LIBRARY_PATH -block. See unslothai/unsloth#5106. +nvidia-cublas-cu13, etc.) and the prebuilt llama-server.exe must find those +DLLs at runtime to load CUDA. Mirrors the Linux LD_LIBRARY_PATH block. +See unslothai/unsloth#5106. """ from __future__ import annotations @@ -60,8 +60,7 @@ from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 def _make_nvidia_layout(prefix: Path, pkgs_with_layout: dict[str, str]): - """Build a fake /Lib/site-packages/nvidia//{bin|Library/bin} - tree with a stub DLL inside each leaf so isdir() picks them up.""" + """Build a fake nvidia//{bin|Library/bin} tree with a stub DLL per leaf.""" nv = prefix / "Lib" / "site-packages" / "nvidia" for pkg, layout in pkgs_with_layout.items(): if layout == "bin": @@ -124,9 +123,8 @@ class TestWindowsPipNvidiaDllDirs: assert len(result) == 4 def test_does_not_walk_outside_known_paths(self, tmp_path): - # Only nvidia//{bin,Library/bin} and torch/lib are picked - # up. Unrelated site-packages contents (numpy, scipy, ...) must - # be ignored. + # Only nvidia//{bin,Library/bin} and torch/lib are picked up. + # Unrelated site-packages contents (numpy, scipy, ...) are ignored. site = tmp_path / "Lib" / "site-packages" (site / "numpy").mkdir(parents = True) (site / "scipy" / "linalg").mkdir(parents = True) @@ -134,10 +132,8 @@ class TestWindowsPipNvidiaDllDirs: assert result == [] def test_picks_up_torch_lib(self, tmp_path): - # PyTorch's Windows CUDA wheel bundles cudart64_X.dll / - # cublas64_X.dll directly under Lib/site-packages/torch/lib/ - # instead of as separate nvidia-* wheels. Without this, users - # on torch-bundled-CUDA installs still hit #5106. + # PyTorch's Windows CUDA wheel bundles cudart64/cublas64 DLLs under + # torch/lib/ rather than as nvidia-* wheels; else still hits #5106. torch_lib = tmp_path / "Lib" / "site-packages" / "torch" / "lib" torch_lib.mkdir(parents = True) (torch_lib / "cudart64_12.dll").write_bytes(b"") @@ -146,8 +142,7 @@ class TestWindowsPipNvidiaDllDirs: assert Path(result[0]) == torch_lib def test_torch_lib_combined_with_nvidia_wheels(self, tmp_path): - # Both modular nvidia-* wheels and torch/lib are returned when - # present together. + # Both modular nvidia-* wheels and torch/lib are returned together. _make_nvidia_layout( tmp_path, { @@ -165,8 +160,7 @@ class TestWindowsPipNvidiaDllDirs: assert any(Path(p) == torch_lib for p in result) def test_torch_lib_must_be_a_directory(self, tmp_path): - # If torch/lib exists as a file (broken install), it is - # ignored, not returned. + # If torch/lib exists as a file (broken install), it is ignored. site = tmp_path / "Lib" / "site-packages" / "torch" site.mkdir(parents = True) (site / "lib").write_bytes(b"not a dir") @@ -176,29 +170,20 @@ class TestWindowsPipNvidiaDllDirs: def test_skips_non_directories(self, tmp_path): nv = tmp_path / "Lib" / "site-packages" / "nvidia" (nv / "cuda_runtime").mkdir(parents = True) - # Create a regular file at the path where 'bin' would normally be a dir + # Regular file where 'bin' would normally be a dir (nv / "cuda_runtime" / "bin").write_bytes(b"not a dir") result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) assert result == [] def test_missing_prefix_does_not_raise(self): - # If sys.prefix points to a path that doesn't exist (unusual, - # but possible during test setup), the resolver must just - # return [] rather than raising. - result = LlamaCppBackend._windows_pip_nvidia_dll_dirs( - "/this/path/does/not/exist/anywhere" - ) + # Nonexistent sys.prefix: resolver must return [], not raise. + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs("/this/path/does/not/exist/anywhere") assert result == [] def test_picks_up_cu13_bin_x86_64_layout(self, tmp_path): - # Current ``nvidia-cuda-runtime`` 13.x and ``nvidia-cublas`` - # 13.x Windows wheels ship DLLs under - # ``nvidia/cu13/bin/x86_64/`` instead of ``nvidia//bin/``. - # Without this, users on the new CUDA 13 wheel generation hit - # the original #5106 failure mode. - dll_dir = ( - tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" - ) + # nvidia 13.x Windows wheels ship DLLs under nvidia/cu13/bin/x86_64/ + # not nvidia//bin/; else the new CUDA 13 wheels hit #5106. + dll_dir = tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" dll_dir.mkdir(parents = True) for name in ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"): (dll_dir / name).write_bytes(b"") @@ -207,7 +192,7 @@ class TestWindowsPipNvidiaDllDirs: def test_picks_up_bin_x64_layout(self, tmp_path): # Some repackaged wheels use ``bin/x64`` (Windows-x64 convention) - # instead of ``bin/x86_64`` (NVIDIA-internal convention). + # rather than ``bin/x86_64`` (NVIDIA-internal convention). dll_dir = tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x64" dll_dir.mkdir(parents = True) (dll_dir / "cudart64_13.dll").write_bytes(b"") @@ -215,9 +200,8 @@ class TestWindowsPipNvidiaDllDirs: assert str(dll_dir) in result def test_mixed_cu12_and_cu13_layouts(self, tmp_path): - # A venv could have both the modular cu12 wheels (legacy) and - # the unsuffixed cu13 wheel installed side by side. Both must - # be reachable. + # A venv could have both the modular cu12 wheels (legacy) and the + # unsuffixed cu13 wheel side by side. Both must be reachable. site = tmp_path / "Lib" / "site-packages" cu12_bin = site / "nvidia" / "cuda_runtime" / "bin" cu13_arch = site / "nvidia" / "cu13" / "bin" / "x86_64" @@ -229,10 +213,8 @@ class TestWindowsPipNvidiaDllDirs: assert cu13_arch in result_set def test_glob_meta_in_prefix_is_safe(self, tmp_path): - # Windows usernames / install paths can contain ``[`` or ``]``. - # A glob-based resolver would interpret these as a character - # class and silently return [] even when DLL dirs exist. The - # iterdir-based implementation must work on such paths. + # Windows paths can contain ``[``/``]``; a glob-based resolver would + # read these as a character class. The iterdir impl must handle them. prefix = tmp_path / "studio_[gpu]_install" dll_dir = prefix / "Lib" / "site-packages" / "nvidia" / "cuda_runtime" / "bin" dll_dir.mkdir(parents = True) @@ -241,18 +223,16 @@ class TestWindowsPipNvidiaDllDirs: assert str(dll_dir) in result, f"bracket-prefixed path returned empty: {result}" def test_arch_subdir_listed_before_parent_bin(self, tmp_path): - # When both ``nvidia//bin/`` and - # ``nvidia//bin/x86_64/`` exist, the arch-specific subdir - # must be listed first so Windows DLL search picks up the - # cudart64_X.dll location even if the parent ``bin`` is empty. + # When both bin/ and bin/x86_64/ exist, the arch subdir must come first + # so the Windows DLL search finds cudart64_X.dll if parent bin is empty. site = tmp_path / "Lib" / "site-packages" outer_bin = site / "nvidia" / "cu13" / "bin" arch_bin = outer_bin / "x86_64" arch_bin.mkdir(parents = True) (arch_bin / "cudart64_13.dll").write_bytes(b"") result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) - # outer_bin exists as a directory (it contains arch_bin); the - # arch-specific subdir should come first in the list. + # outer_bin exists as a dir (it holds arch_bin); the arch-specific + # subdir should come first in the list. result_paths = [Path(p) for p in result] assert arch_bin in result_paths assert outer_bin in result_paths diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py new file mode 100644 index 0000000000..bf0c4b731f --- /dev/null +++ b/studio/backend/tests/test_llama_route.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""routes/llama.py: the source_build field is exposed and the handlers run the +(now subprocess-touching) detection off the event loop via a worker thread. + +The route file is loaded standalone with a stubbed auth dependency so the test +does not pull the whole routes package (matplotlib-heavy training router) and +works in a minimal env. +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +import threading +import types +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +pytest.importorskip("fastapi") + + +def _load_route(): + # Prefer the real auth module; stub it only in minimal envs where its + # deps are absent. Stubs are popped after the load so they never leak + # into sys.modules for the rest of the suite. + stubbed = [] + try: + import auth.authentication # noqa: F401 + except Exception: + auth_pkg = types.ModuleType("auth") + auth_pkg.__path__ = [] + auth_mod = types.ModuleType("auth.authentication") + auth_mod.get_current_subject = lambda: "test" + for name, stub in (("auth", auth_pkg), ("auth.authentication", auth_mod)): + if name not in sys.modules: + sys.modules[name] = stub + stubbed.append(name) + try: + spec = importlib.util.spec_from_file_location( + "llama_route_under_test", str(_BACKEND / "routes" / "llama.py") + ) + mod = importlib.util.module_from_spec(spec) + sys.modules["llama_route_under_test"] = mod # so pydantic resolves forward refs + spec.loader.exec_module(mod) + return mod + finally: + for name in stubbed: + sys.modules.pop(name, None) + + +rl = _load_route() + + +def test_status_response_exposes_source_build(): + payload = { + "supported": True, + "update_available": True, + "stale": False, + "installed_tag": None, + "latest_tag": "b9585", + "published_repo": "unslothai/llama.cpp", + "installed_at_utc": None, + "age_days": None, + "source_build": True, + "job": {"state": "idle"}, + } + model = rl.LlamaUpdateStatusResponse(**payload) + assert model.model_dump()["source_build"] is True + # Extra/unknown keys must not crash the response model. + rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1}) + + +def test_status_handler_runs_off_event_loop(monkeypatch): + seen = {} + + def fake_status(force_refresh = False): + seen["thread"] = threading.current_thread() + return { + "supported": True, + "update_available": True, + "source_build": True, + "latest_tag": "b9585", + "job": {"state": "idle"}, + } + + monkeypatch.setattr(rl, "get_update_status", fake_status) + out = asyncio.run(rl.llama_update_status(force_refresh = False, current_subject = "t")) + assert out.source_build is True + # Detection ran in a worker thread, not the event-loop thread. + assert seen["thread"] is not threading.main_thread() + + +def test_update_handler_runs_off_event_loop(monkeypatch): + seen = {} + + def fake_start(): + seen["thread"] = threading.current_thread() + return {"started": True, "reason": None, "job": {"state": "running"}} + + monkeypatch.setattr(rl, "start_update", fake_start) + out = asyncio.run(rl.llama_update(current_subject = "t")) + assert out.started is True + assert seen["thread"] is not threading.main_thread() diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 02a272ba3e..09775707ee 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -4,8 +4,8 @@ """Unit tests for the llama-server pass-through args validator. The validator is the boundary between user CLI/HTTP input and the -llama-server subprocess. These tests pin denylist behaviour so it -doesn't quietly regress when new managed flags are added. +llama-server subprocess. These tests pin denylist behaviour so it doesn't +regress when new managed flags are added. """ from __future__ import annotations @@ -16,20 +16,16 @@ from pathlib import Path import pytest -# Load llama_server_args.py directly so this test doesn't drag in the -# full backend chain (fastapi / structlog / loggers / utils.hardware) -# via core/inference/__init__.py. The validator is intentionally -# dependency-free and unit-tests should reflect that. -_LSA_PATH = ( - Path(__file__).resolve().parent.parent - / "core" - / "inference" - / "llama_server_args.py" -) +# Load llama_server_args.py directly to avoid dragging in the full backend +# chain via core/inference/__init__.py. The validator is dependency-free. +_LSA_PATH = Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_server_args.py" _spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH) _lsa = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_lsa) is_managed_flag = _lsa.is_managed_flag +parse_cache_override = _lsa.parse_cache_override +parse_ctx_override = _lsa.parse_ctx_override +resolve_cache_type_kv = _lsa.resolve_cache_type_kv strip_shadowing_flags = _lsa.strip_shadowing_flags validate_extra_args = _lsa.validate_extra_args @@ -74,10 +70,9 @@ validate_extra_args = _lsa.validate_extra_args # Reasoning controls ["--reasoning-format", "deepseek"], ["-rea", "auto"], - # Soft-managed: user-supplied flags last-wins-override Studio's - # auto-set version. --parallel / -np / --n-parallel are NOT - # here -- they're hard-denied (KV-cache + slot count would - # desync). Use `unsloth studio run --parallel N` instead. + # Soft-managed: user flags last-wins over Studio's auto-set version. + # --parallel / -np / --n-parallel are hard-denied (KV-cache + slot + # count would desync); use `unsloth studio run --parallel N` instead. ["-c", "131072"], ["--ctx-size", "8192"], ["--flash-attn", "off"], @@ -126,8 +121,8 @@ def test_non_flag_token_passes_through(): "-np", "--parallel", "--n-parallel", - # Model identity (every alias; bumping llama.cpp must keep - # every form rejected, not just the long). + # Model identity (every alias; bumping llama.cpp must keep every + # form rejected, not just the long one). "-m", "--model", "-mu", @@ -175,8 +170,8 @@ def test_non_flag_token_passes_through(): "--models-max", "--models-autoload", "--no-models-autoload", - # Server-mode flips: --embedding / --rerank would restrict - # llama-server to those endpoints and break Studio's chat hop. + # Server-mode flips: --embedding / --rerank restrict llama-server to + # those endpoints and break Studio's chat hop. "--embedding", "--embeddings", "--rerank", @@ -193,16 +188,16 @@ def test_denylist_rejects_all_aliases(denied): @pytest.mark.parametrize( "args,offending", [ - # Pass-through --parallel would last-wins-override the real - # slot count while Studio's KV-cache fit + llama_parallel_slots - # stay at the typer value -- plan vs. process disagree. + # Pass-through --parallel would last-wins-override the real slot + # count while Studio's KV-cache fit + llama_parallel_slots stay at + # the typer value -- plan vs. process disagree. (["--parallel", "8"], "--parallel"), (["--parallel=8"], "--parallel"), (["--n-parallel", "16"], "--n-parallel"), (["--n-parallel=16"], "--n-parallel"), (["-np", "32"], "-np"), - # Attached short form: Click clusters it CLI-side; HTTP /load - # with `["-np8"]` must still resolve to managed. + # Attached short form: Click clusters it CLI-side; HTTP /load with + # `["-np8"]` must still resolve to managed. (["-np8"], "-np"), (["-np64"], "-np"), # Out-of-range values that would bypass the typer 1..64 guard. @@ -229,8 +224,8 @@ def test_denylist_rejects_equals_form(): [" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"], ) def test_denylist_rejects_whitespace_padded_forms(padded): - # `_flag_name` trims whitespace before lookup; otherwise a trailing - # space could slip a managed flag past the boundary. + # `_flag_name` trims whitespace before lookup; else a trailing space + # could slip a managed flag past the boundary. with pytest.raises(ValueError, match = "parallel|np"): validate_extra_args([padded, "8"]) @@ -240,15 +235,15 @@ def test_denylist_rejects_whitespace_padded_forms(padded): ["-np8x", "-np-1foo", "-np+1bar", "-np9zzz"], ) def test_denylist_rejects_np_with_digit_prefix_and_junk(attached): - # Backend `_flag_name` must classify the same forms the CLI - # rewriter expands, else HTTP /load could smuggle `-np8x` through. + # Backend `_flag_name` must classify the same forms the CLI rewriter + # expands, else HTTP /load could smuggle `-np8x` through. with pytest.raises(ValueError, match = "np"): validate_extra_args([attached]) def test_denylist_rejects_short_form_when_long_is_denied(): - # `-m` is the short form of --model; rejecting only the long - # form would leave a trivial bypass. + # `-m` is the short form of --model; rejecting only the long form + # would leave a trivial bypass. with pytest.raises(ValueError, match = "-m"): validate_extra_args(["-m", "/some/other/path.gguf"]) @@ -356,14 +351,7 @@ def test_strip_shadowing_flags_keeps_cache_when_cache_disabled(): ["--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", - ] + 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(): @@ -371,14 +359,7 @@ def test_strip_shadowing_flags_keeps_spec_when_spec_disabled(): ["--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", - ] + assert out == ["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"] def test_strip_shadowing_flags_drops_mtp_flags_when_requested(): @@ -410,6 +391,86 @@ def test_is_managed_flag_false_for_mtp_pass_through(): assert is_managed_flag("--spec-ngram-mod-n-max") is False +# ── parse_ctx_override ─────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "args,expected", + [ + (None, None), + ([], None), + (["--top-k", "20"], None), + (["--ctx-size", "128000"], 128000), + (["--ctx-size=128000"], 128000), + (["-c", "128000"], 128000), + (["-c=128000"], 128000), + (["-c", "4096", "--ctx-size", "128000"], 128000), + ], +) +def test_parse_ctx_override(args, expected): + assert parse_ctx_override(args) == expected + + +@pytest.mark.parametrize( + "args", + [ + ["--ctx-size"], + ["--ctx-size", "--top-k"], + ["--ctx-size", "abc"], + ["--ctx-size=abc"], + ["-c", "-1"], + ], +) +def test_parse_ctx_override_rejects_malformed_values(args): + with pytest.raises(ValueError, match = "ctx-size|'-c'"): + parse_ctx_override(args) + + +def test_validate_extra_args_rejects_malformed_ctx_override(): + with pytest.raises(ValueError, match = "ctx-size"): + validate_extra_args(["--ctx-size", "abc"]) + + +# ── parse_cache_override ───────────────────────────────────────────── + + +@pytest.mark.parametrize( + "args,expected", + [ + (None, None), + ([], None), + (["--top-k", "20"], None), + (["--cache-type-k", "q8_0"], "q8_0"), + (["-ctk", "q4_0"], "q4_0"), + (["-ctv", "q4_0"], "q4_0"), + (["--cache-type-k=q4_0"], "q4_0"), + (["-ctk", "f16", "-ctk", "q8_0"], "q8_0"), + ], +) +def test_parse_cache_override(args, expected): + assert parse_cache_override(args) == expected + + +@pytest.mark.parametrize( + "args", + [ + ["-ctk"], + ["-ctk", "-c", "4096"], + ], +) +def test_parse_cache_override_rejects_malformed_values(args): + with pytest.raises(ValueError, match = "cache-type|'-ctk'"): + parse_cache_override(args) + + +def test_resolve_cache_type_kv_uses_override_when_present(): + assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0" + + +def test_resolve_cache_type_kv_uses_fallback_without_override(): + assert resolve_cache_type_kv(["--top-k", "20"], "f16") == "f16" + + def test_strip_shadowing_flags_boolean_does_not_consume_next_token(): # `--spec-default` is boolean; drop just the flag, keep the next token. out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True) @@ -422,9 +483,7 @@ def test_strip_shadowing_flags_jinja_boolean_preserves_positional(): def test_strip_shadowing_flags_no_jinja_boolean_preserves_positional(): - out = strip_shadowing_flags( - ["--no-jinja", "trailing-positional"], strip_template = True - ) + out = strip_shadowing_flags(["--no-jinja", "trailing-positional"], strip_template = True) assert out == ["trailing-positional"] @@ -448,3 +507,28 @@ def test_strip_shadowing_flags_defaults_strip_everything(): ["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"] ) assert out == [] + + +def test_strip_shadowing_flags_drops_model_draft_with_spec(): + # --model-draft (and aliases) are Studio-managed since the separate + # MTP drafter support: an inherited copy must not last-wins-override + # the auto-detected drafter. + out = strip_shadowing_flags( + ["--model-draft", "/old/mtp.gguf", "-md", "/old2.gguf", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = True, + strip_template = False, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_model_draft_without_spec(): + out = strip_shadowing_flags( + ["--model-draft", "/custom/mtp.gguf"], + strip_context = True, + strip_cache = False, + strip_spec = False, + strip_template = False, + ) + assert out == ["--model-draft", "/custom/mtp.gguf"] diff --git a/studio/backend/tests/test_llm_assist_startup_opt_in.py b/studio/backend/tests/test_llm_assist_startup_opt_in.py new file mode 100644 index 0000000000..e81b1d3775 --- /dev/null +++ b/studio/backend/tests/test_llm_assist_startup_opt_in.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for Helper LLM startup pre-cache opt-in behavior.""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from models.datasets import AiAssistMappingRequest +from routes import datasets as datasets_route +from routes import settings as settings_route +from utils import helper_precache_settings + + +def _install_fake_studio_db(monkeypatch, *, stored = None): + storage_pkg = types.ModuleType("storage") + studio_db = types.ModuleType("storage.studio_db") + values: dict[str, object] = {} + if stored is not None: + values[helper_precache_settings.HELPER_PRECACHE_SETTING_KEY] = stored + + def get_app_setting(key, fallback = None): + return values.get(key, fallback) + + def upsert_app_settings(settings): + values.update(settings) + return dict(values) + + studio_db.get_app_setting = get_app_setting + studio_db.upsert_app_settings = upsert_app_settings + monkeypatch.setitem(sys.modules, "storage", storage_pkg) + monkeypatch.setitem(sys.modules, "storage.studio_db", studio_db) + return values + + +def test_helper_precache_defaults_off_when_setting_missing(monkeypatch): + monkeypatch.delenv("UNSLOTH_HELPER_MODEL_DISABLE", raising = False) + _install_fake_studio_db(monkeypatch) + + assert helper_precache_settings.get_helper_precache_enabled() is False + assert helper_precache_settings.should_preload_helper_on_startup() is False + + +def test_helper_precache_opt_in_is_blocked_by_existing_disable_env(monkeypatch): + _install_fake_studio_db(monkeypatch, stored = True) + monkeypatch.setenv("UNSLOTH_HELPER_MODEL_DISABLE", "true") + + assert helper_precache_settings.get_helper_precache_enabled() is True + assert helper_precache_settings.should_preload_helper_on_startup() is False + + +def test_settings_route_persists_helper_precache_toggle(monkeypatch): + values = _install_fake_studio_db(monkeypatch) + monkeypatch.delenv("UNSLOTH_HELPER_MODEL_DISABLE", raising = False) + + response = settings_route.update_helper_precache( + settings_route.HelperPrecachePayload(enabled = True), + current_subject = "test-user", + ) + + assert response.enabled is True + assert response.default_enabled is False + assert response.disabled_by_env is False + assert values[helper_precache_settings.HELPER_PRECACHE_SETTING_KEY] is True + + +def test_main_startup_uses_helper_precache_gate_instead_of_unconditional_precache(): + source = (Path(__file__).resolve().parent.parent / "main.py").read_text(encoding = "utf-8") + startup_section = source[ + source.index("cleanup_orphaned_runs") : source.index("# Initialize RSA key pair") + ] + + assert "_start_helper_precache_if_enabled()" in startup_section + assert "precache_helper_gguf" not in startup_section + assert "threading.Thread(target = _precache" not in startup_section + + +def test_ai_assist_route_still_calls_on_demand_advisor(monkeypatch): + calls: list[dict] = [] + llm_assist = types.ModuleType("utils.datasets.llm_assist") + + def fake_llm_conversion_advisor(**kwargs): + calls.append(kwargs) + return { + "success": True, + "suggested_mapping": {"prompt": "user", "answer": "assistant"}, + "system_prompt": "Answer carefully.", + "dataset_type": "question_answering", + "is_conversational": False, + "user_notification": "Columns mapped by AI Assist.", + } + + llm_assist.llm_conversion_advisor = fake_llm_conversion_advisor + monkeypatch.setitem(sys.modules, "utils.datasets.llm_assist", llm_assist) + + response = datasets_route.ai_assist_mapping( + AiAssistMappingRequest( + columns = ["prompt", "answer"], + samples = [{"prompt": "x" * 250, "answer": "ok", "extra": "ignored"}], + dataset_name = "owner/dataset", + hf_token = "hf_test", + model_name = "unsloth/test", + model_type = "text", + ), + current_subject = "test-user", + ) + + assert response.success is True + assert response.suggested_mapping == {"prompt": "user", "answer": "assistant"} + assert response.system_prompt == "Answer carefully." + assert calls == [ + { + "column_names": ["prompt", "answer"], + "samples": [{"prompt": "x" * 200, "answer": "ok"}], + "dataset_name": "owner/dataset", + "hf_token": "hf_test", + "model_name": "unsloth/test", + "model_type": "text", + } + ] diff --git a/studio/backend/tests/test_log_filter_no_truncation.py b/studio/backend/tests/test_log_filter_no_truncation.py index d78643f5b9..52ffa2ba32 100644 --- a/studio/backend/tests/test_log_filter_no_truncation.py +++ b/studio/backend/tests/test_log_filter_no_truncation.py @@ -2,27 +2,11 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Regression tests for studio.backend.loggers.handlers.filter_sensitive_data. +Regression tests for loggers.handlers.filter_sensitive_data. -Context: filter_sensitive_data was originally written with a base64-detection -heuristic that truncated any string >100 chars containing ',' or '/' down to -20 chars + '...'. The block was dormant until PR #5246 wired the processor -into the structlog chain to redact native-path leases. Once active, the -heuristic ate normal log lines emitted by llama_cpp_backend (GGUF size -summary, mmproj selection, the full llama-server command line) and any -exception traceback that happened to contain a file path. - -These tests pin two properties: - -1. Long, comma- or slash-bearing log messages flow through filter_sensitive_data - unchanged. The exact strings exercised match the call sites at - studio/backend/core/inference/llama_cpp.py:2117, :2283, and :2312 that - were truncated in the original bug report. - -2. PR #5246's native-path lease redaction still fires for both the inline - ``native_path_lease=...`` regex form and the ``nativePathLease`` dict-key - form. This guards against future regressions that strip redaction along - with the truncation block. +Pins two properties: (1) long strings with commas/slashes pass through unchanged +(the base64-truncation heuristic from PR #5246 was too aggressive), and +(2) native-path lease redaction still fires for inline and dict-key forms. """ from loggers.handlers import filter_sensitive_data diff --git a/studio/backend/tests/test_login_rate_limit.py b/studio/backend/tests/test_login_rate_limit.py index c8498d4857..14b10576da 100644 --- a/studio/backend/tests/test_login_rate_limit.py +++ b/studio/backend/tests/test_login_rate_limit.py @@ -4,11 +4,11 @@ """Tests for the per-(ip, username) login rate limiter. Covers: - - bucket key composition is (client-ip, username.lower()) - - X-Forwarded-For is honoured only when UNSLOTH_STUDIO_TRUST_FORWARDED is set + - bucket key is (client-ip, username.lower()) + - X-Forwarded-For honoured only when UNSLOTH_STUDIO_TRUST_FORWARDED is set - 429 detail body does NOT leak the client IP - - One username failing does not lock out a different user from the same IP - - One IP failing does not lock out the same user from a different IP + - One username failing doesn't lock out a different user from the same IP + - One IP failing doesn't lock out the same user from a different IP """ import os @@ -45,9 +45,12 @@ def env_trust_proxy(monkeypatch): class _FakeRequest: - def __init__(self, client_host = "127.0.0.1", headers = None): + def __init__( + self, + client_host = "127.0.0.1", + headers = None, + ): from starlette.datastructures import Headers - self.client = type("Client", (), {"host": client_host})() self.headers = Headers(headers or {}) @@ -58,23 +61,19 @@ class _FakeRequest: class TestClientIp: def test_uses_request_client_host_by_default(self, env_no_proxy): from routes.auth import _client_ip - assert _client_ip(_FakeRequest("203.0.113.5")) == "203.0.113.5" def test_ignores_xff_when_trust_off(self, env_no_proxy): from routes.auth import _client_ip - req = _FakeRequest( "127.0.0.1", {"x-forwarded-for": "198.51.100.7, 10.0.0.1"}, ) - # The proxy header could be spoofed; without the opt-in we - # only trust the direct connection. + # Proxy header is spoofable; without the opt-in, trust the direct connection. assert _client_ip(req) == "127.0.0.1" def test_honours_first_xff_when_trust_on(self, env_trust_proxy): from routes.auth import _client_ip - req = _FakeRequest( "127.0.0.1", {"x-forwarded-for": "198.51.100.7, 10.0.0.1"}, @@ -83,12 +82,10 @@ class TestClientIp: def test_falls_back_to_client_host_when_xff_missing(self, env_trust_proxy): from routes.auth import _client_ip - assert _client_ip(_FakeRequest("203.0.113.9")) == "203.0.113.9" def test_honours_forwarded_header_when_trust_on(self, env_trust_proxy): from routes.auth import _client_ip - req = _FakeRequest( "127.0.0.1", {"forwarded": 'for="198.51.100.42";proto=https'}, @@ -104,41 +101,29 @@ class TestClientIp: def test_xff_strips_ipv4_port(self, env_trust_proxy): from routes.auth import _client_ip - - req = _FakeRequest( - "127.0.0.1", {"x-forwarded-for": "198.51.100.7:50001, 10.0.0.1"} - ) + req = _FakeRequest("127.0.0.1", {"x-forwarded-for": "198.51.100.7:50001, 10.0.0.1"}) assert _client_ip(req) == "198.51.100.7" def test_xff_strips_bracketed_ipv6_port(self, env_trust_proxy): from routes.auth import _client_ip - - req = _FakeRequest( - "127.0.0.1", {"x-forwarded-for": "[2001:db8::1]:50001, 10.0.0.1"} - ) + req = _FakeRequest("127.0.0.1", {"x-forwarded-for": "[2001:db8::1]:50001, 10.0.0.1"}) assert _client_ip(req) == "2001:db8::1" def test_forwarded_strips_ipv4_port(self, env_trust_proxy): from routes.auth import _client_ip - - req = _FakeRequest( - "127.0.0.1", {"forwarded": 'for="198.51.100.7:50001";proto=https'} - ) + req = _FakeRequest("127.0.0.1", {"forwarded": 'for="198.51.100.7:50001";proto=https'}) assert _client_ip(req) == "198.51.100.7" def test_forwarded_strips_bracketed_ipv6_port(self, env_trust_proxy): from routes.auth import _client_ip - - req = _FakeRequest( - "127.0.0.1", {"forwarded": 'for="[2001:db8::1]:50001";proto=https'} - ) + req = _FakeRequest("127.0.0.1", {"forwarded": 'for="[2001:db8::1]:50001";proto=https'}) assert _client_ip(req) == "2001:db8::1" def test_forwarded_isolates_first_element(self, env_trust_proxy): from routes.auth import _client_ip - # Multi-element Forwarded must pick the first element only, - # otherwise suffix variations create attacker-controlled buckets. + # Pick the first Forwarded element only, else suffix variations create + # attacker-controlled buckets. req = _FakeRequest( "127.0.0.1", {"forwarded": "for=198.51.100.42, for=10.0.0.1;proto=https"}, @@ -169,7 +154,7 @@ class TestBucketKeyAndBlocking: for _ in range(_LOGIN_MAX_FAILS): _record_login_failure(_bucket_key(req, "alice")) assert _login_blocked(_bucket_key(req, "alice")) > 0 - # bob's account from the same IP is unaffected by alice's typos. + # bob's account from the same IP is unaffected by alice's typos assert _login_blocked(_bucket_key(req, "bob")) == 0 def test_record_per_ip_isolates_other_ips(self, env_no_proxy): @@ -203,9 +188,7 @@ class TestBucketKeyAndBlocking: req = _FakeRequest("203.0.113.10") for idx in range(5): auth_routes._record_login_failure(auth_routes._unknown_user_key(req)) - # Different "username" each attempt would not have throttled - # under per-(ip,username) only; the IP aggregate must. - # The next missing-user attempt is blocked. + # Per-(ip,username) alone wouldn't throttle distinct usernames; the IP aggregate must. assert auth_routes._login_blocked(auth_routes._unknown_user_key(req)) > 0 def test_unknown_user_bucket_is_single_sentinel(self, env_no_proxy): @@ -216,8 +199,7 @@ class TestBucketKeyAndBlocking: unknown_key = auth_routes._unknown_user_key(req) for _ in range(20): auth_routes._record_login_failure(unknown_key) - # Account bucket cardinality stays at exactly one sentinel entry - # for this IP regardless of how many distinct usernames sprayed. + # Exactly one sentinel bucket for this IP regardless of usernames sprayed. ip_keys = [k for k in auth_routes._LOGIN_BUCKETS if k[0] == "203.0.113.11"] assert len(ip_keys) == 1 assert ip_keys[0][1].startswith("\x00") @@ -230,7 +212,7 @@ class TestBucketKeyAndBlocking: req = _FakeRequest("203.0.113.12") for idx in range(50): auth_routes._record_login_failure((req.client.host, f"user-{idx}")) - # Hard cap respected; further keys do not allocate. + # Hard cap respected; further keys don't allocate. assert len(auth_routes._LOGIN_BUCKETS) <= 10 @@ -247,9 +229,7 @@ class TestLogin429Body: import secrets as _secrets monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr( - storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password" - ) + monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password") monkeypatch.setattr(storage, "_bootstrap_password", None) storage.create_initial_user( username = storage.DEFAULT_ADMIN_USERNAME, @@ -265,7 +245,7 @@ class TestLogin429Body: def test_429_detail_does_not_leak_ip(self, env_no_proxy, login_client): from routes.auth import _LOGIN_MAX_FAILS - # Drive 6 failures from the same client IP / username. + # Drive 6 failures from the same client IP / username for _ in range(_LOGIN_MAX_FAILS): r = login_client.post( "/api/auth/login", @@ -278,8 +258,8 @@ class TestLogin429Body: ) assert r.status_code == 429 detail = r.json()["detail"] - # The 429 body must not interpolate the source IP. + # The 429 body must not interpolate the source IP assert "127.0.0.1" not in detail assert "Too many" in detail - # Retry-After header is still set for clients. + # Retry-After header is still set for clients assert "Retry-After" in r.headers diff --git a/studio/backend/tests/test_mcp_config_import.py b/studio/backend/tests/test_mcp_config_import.py new file mode 100644 index 0000000000..4082733490 --- /dev/null +++ b/studio/backend/tests/test_mcp_config_import.py @@ -0,0 +1,341 @@ +"""Tests for MCP config-file import (issue #5936). + +Covers the round-trip-safe command join/split inverse (join_stdio_command ↔ +parse_stdio_command, on both posix and win32 using the issue's Windows +fixtures), the pure config parser (parse_mcp_config), and the POST /import +route (stdio gate on/off, url dedup, one bad entry not sinking the batch). + +Run from studio/backend: python -m pytest tests/test_mcp_config_import.py -q +""" + +import sys + +import pytest + +from core.inference import mcp_client +from core.inference.mcp_config_import import parse_mcp_config +from storage import mcp_servers_db + + +def _reset_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + + +def _enable(monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + + +def _disable(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False) + + +# ── 1. join_stdio_command ↔ parse_stdio_command round-trip ────────── + + +@pytest.mark.parametrize( + "parts", + [ + ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp"], + ["python", "-m", "mod", "--name", "a b"], + ["uvx", "some-server", "--flag"], + ["/usr/local/bin/my-server"], + ["mcp-server-sqlite"], + ], +) +def test_join_parse_roundtrip_posix(monkeypatch, parts): + monkeypatch.setattr(sys, "platform", "linux") + joined = mcp_client.join_stdio_command(parts) + assert mcp_client.parse_stdio_command(joined) == parts + + +@pytest.mark.parametrize( + "parts", + [ + # Issue #5936's literal Windows examples: absolute .exe with a path, and + # backslash drive/dir args must survive the join→split round-trip intact. + [ + "C:\\Users\\user\\Documents\\Office-Word-MCP-Server\\.venv\\Scripts\\python.exe", + "C:\\Users\\user\\Documents\\Office-Word-MCP-Server\\word_mcp_server.py", + ], + [ + "node", + "C:\\Users\\user\\Documents\\DesktopCommanderMCP\\dist\\index.js", + "--no-onboarding", + ], + [ + "node", + "C:\\Users\\user\\AppData\\Roaming\\npm\\node_modules\\@modelcontextprotocol\\server-filesystem\\dist\\index.js", + "D:\\", + "O:\\", + ], + # A command path with spaces is the case that actually needs quoting. + ["C:\\Program Files\\node\\node.exe", "server.js"], + ["C:\\Program Files\\Foo\\", "server.js"], + ["C:\\Program Files\\Foo\\", '{"foo":"bar"}'], + ["'C:\\Program Files\\node\\node.exe'", "server.js"], + ["node", "O'Reilly"], + ["node", "C:\\Users\\O'Reilly\\server.js"], + ["node", "'draft'"], + ["node", "'open", "close'"], + ["node", ""], + ], +) +def test_join_parse_roundtrip_win32(monkeypatch, parts): + monkeypatch.setattr(sys, "platform", "win32") + joined = mcp_client.join_stdio_command(parts) + assert mcp_client.parse_stdio_command(joined) == parts + + +def test_parse_rejects_manual_single_quoted_windows_executable(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + command = "'C:\\Program Files\\node\\node.exe' server.js" + with pytest.raises(ValueError): + mcp_client.parse_stdio_command(command) + + +def test_parse_windows_apostrophes_as_literals(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + assert mcp_client.parse_stdio_command("node O'Reilly") == ["node", "O'Reilly"] + assert mcp_client.parse_stdio_command("node C:\\Users\\O'Reilly\\server.js") == [ + "node", + "C:\\Users\\O'Reilly\\server.js", + ] + assert mcp_client.parse_stdio_command("node 'draft'") == ["node", "'draft'"] + assert mcp_client.parse_stdio_command("node 'open close'") == ["node", "'open", "close'"] + + +def test_parse_rejects_unterminated_windows_double_quote(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + with pytest.raises(ValueError): + mcp_client.parse_stdio_command('node "C:\\path with spaces') + + +# ── 2. parse_mcp_config ───────────────────────────────────────────── + + +def test_parse_stdio_entry(): + cfg = { + "mcpServers": { + "fs": { + "command": "npx", + "args": ["-y", "server", "/tmp"], + "env": {"K": "v"}, + } + } + } + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert len(entries) == 1 + entry = entries[0] + assert entry.display_name == "fs" + assert entry.is_stdio is True + assert entry.headers == {"K": "v"} + assert mcp_client.parse_stdio_command(entry.url) == ["npx", "-y", "server", "/tmp"] + + +def test_parse_remote_entry(): + cfg = { + "mcpServers": { + "remote": { + "url": "https://example.com/mcp", + "headers": {"Authorization": "Bearer x"}, + } + } + } + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].url == "https://example.com/mcp" + assert entries[0].is_stdio is False + assert entries[0].headers == {"Authorization": "Bearer x"} + + +def test_parse_preserves_disabled_and_oauth(): + cfg = { + "servers": { + "remote": { + "type": "http", + "url": "https://example.com/mcp", + "oauth": {"clientId": "client"}, + "disabled": True, + } + } + } + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].is_enabled is False + assert entries[0].use_oauth is True + + +def test_parse_accepts_cline_streamable_http_alias(): + cfg = { + "mcpServers": { + "remote": { + "type": "streamableHttp", + "url": "https://example.com/mcp", + } + } + } + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].url == "https://example.com/mcp" + assert entries[0].is_stdio is False + + +@pytest.mark.parametrize( + "server", + [ + {"command": "node", "args": ["server.js"], "cwd": "/tmp/server"}, + {"command": "node", "args": ["server.js"], "envFile": ".env"}, + {"command": "node", "args": ["server.js"], "env": {"API_KEY": "${input:api-key}"}}, + {"command": "node", "args": ["${workspaceFolder}/server.js"]}, + {"command": "node", "args": ["server.js"], "env": {"HTTP_PROXY": None}}, + {"command": "node", "args": ["server.js"], "sandboxEnabled": True}, + {"url": "https://example.com/mcp", "headers": {"Authorization": "Bearer ${input:token}"}}, + {"url": "https://example.com/mcp", "headers": {"Authorization": None}}, + {"type": "http", "url": "https://example.com/sse"}, + {"type": "http", "url": "https://example.com/sse "}, + {"type": "streamableHttp", "url": "https://example.com/sse"}, + {"url": "https://example.com/mcp", "timeout": 120}, + {"url": "https://example.com/mcp", "timeoutMs": 120000}, + {"url": "https://example.com/mcp", "timeoutSeconds": 120}, + {"type": "sse", "url": "https://example.com/custom"}, + ], +) +def test_parse_rejects_unrepresentable_imports(server): + entries, errors = parse_mcp_config({"servers": {"bad": server}}) + assert entries == [] + assert len(errors) == 1 + + +def test_servers_alias_key(): + # VS Code uses "servers" instead of "mcpServers". + cfg = {"servers": {"fs": {"command": "node", "args": ["x.js"]}}} + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert len(entries) == 1 + + +def test_env_and_args_values_coerced_to_str(): + cfg = {"mcpServers": {"fs": {"command": "node", "args": [8080], "env": {"PORT": 8080}}}} + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].headers == {"PORT": "8080"} + assert mcp_client.parse_stdio_command(entries[0].url) == ["node", "8080"] + + +def test_args_optional(): + cfg = {"mcpServers": {"sqlite": {"command": "mcp-server-sqlite"}}} + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].url == "mcp-server-sqlite" + assert entries[0].headers is None + + +def test_bad_entry_does_not_sink_batch(): + cfg = { + "mcpServers": { + "good": {"command": "node", "args": ["x.js"]}, + "both": {"command": "node", "url": "https://x/mcp"}, + "neither": {"name": "oops"}, + "bad_args": {"command": "node", "args": "x.js"}, + "bad_env": {"command": "node", "env": ["NOT", "A", "DICT"]}, + } + } + entries, errors = parse_mcp_config(cfg) + assert {e.display_name for e in entries} == {"good"} + assert len(errors) == 4 + + +def test_not_a_dict(): + entries, errors = parse_mcp_config([]) + assert entries == [] + assert len(errors) == 1 + + +def test_missing_servers_key(): + entries, errors = parse_mcp_config({"foo": {}}) + assert entries == [] + assert len(errors) == 1 + + +def test_servers_alias_error_names_actual_key(): + entries, errors = parse_mcp_config({"servers": []}) + assert entries == [] + assert errors == ["'servers' must be an object mapping name -> server."] + + +# ── 3. POST /import route ─────────────────────────────────────────── + + +def test_import_route_creates_and_dedups(tmp_path, monkeypatch): + import asyncio + + from models.mcp_servers import McpServerImportRequest + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + cfg = { + "mcpServers": { + "fs": { + "command": "npx", + "args": ["-y", "server", "/tmp"], + "env": {"API_KEY": "sk"}, + }, + "remote": {"url": "https://example.com/mcp"}, + "oauth": { + "type": "http", + "url": "https://auth.example.com/mcp", + "oauth": {"clientId": "client"}, + }, + "disabled": { + "url": "https://disabled.example.com/mcp", + "disabled": True, + }, + } + } + res = asyncio.run( + routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + ) + assert res.errors == [] + assert res.skipped == [] + assert {c.display_name for c in res.created} == {"fs", "remote", "oauth", "disabled"} + fs = next(c for c in res.created if c.display_name == "fs") + assert fs.headers == {"API_KEY": "sk"} + assert fs.use_oauth is False + assert fs.is_enabled is True + oauth = next(c for c in res.created if c.display_name == "oauth") + assert oauth.use_oauth is True + disabled = next(c for c in res.created if c.display_name == "disabled") + assert disabled.is_enabled is False + + # Re-importing the same config skips both by url. + res2 = asyncio.run( + routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + ) + assert res2.created == [] + assert set(res2.skipped) == {"fs", "remote", "oauth", "disabled"} + + +def test_import_route_gates_stdio_when_disabled(tmp_path, monkeypatch): + import asyncio + + from models.mcp_servers import McpServerImportRequest + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + _disable(monkeypatch) + cfg = { + "mcpServers": { + "fs": {"command": "npx", "args": ["server"]}, + "remote": {"url": "https://example.com/mcp"}, + } + } + res = asyncio.run( + routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + ) + # Remote still imports; the stdio entry is rejected per-entry (gate off). + assert {c.display_name for c in res.created} == {"remote"} + assert any("fs" in err for err in res.errors) + assert len(mcp_servers_db.list_servers()) == 1 diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py new file mode 100644 index 0000000000..ede3cf15d4 --- /dev/null +++ b/studio/backend/tests/test_mcp_servers.py @@ -0,0 +1,600 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import pytest +from fastapi import HTTPException + +from storage import mcp_servers_db + + +def _reset_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + + +# ── storage: mcp_servers_db ───────────────────────────────────────── + + +def test_create_and_get_server(tmp_path, monkeypatch): + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server( + id = "srv1", + display_name = "GitHub", + url = "https://example.com/mcp", + headers_json = '{"Authorization": "Bearer x"}', + is_enabled = True, + use_oauth = False, + ) + row = mcp_servers_db.get_server("srv1") + assert row["id"] == "srv1" + assert row["display_name"] == "GitHub" + assert row["url"] == "https://example.com/mcp" + assert row["headers_json"] == '{"Authorization": "Bearer x"}' + assert row["is_enabled"] == 1 + assert row["use_oauth"] == 0 + + +def test_list_servers_ordered_by_created_at(tmp_path, monkeypatch): + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server(id = "a", display_name = "A", url = "https://a/m") + mcp_servers_db.create_server(id = "b", display_name = "B", url = "https://b/m") + rows = mcp_servers_db.list_servers() + assert [r["id"] for r in rows] == ["a", "b"] + + +def test_update_server_coerces_bools(tmp_path, monkeypatch): + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m") + assert mcp_servers_db.update_server("srv1", {"is_enabled": False, "use_oauth": True}) + row = mcp_servers_db.get_server("srv1") + assert row["is_enabled"] == 0 + assert row["use_oauth"] == 1 + + +def test_update_server_empty_changes_returns_false(tmp_path, monkeypatch): + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m") + assert mcp_servers_db.update_server("srv1", {}) is False + + +def test_delete_server_roundtrip(tmp_path, monkeypatch): + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m") + assert mcp_servers_db.delete_server("srv1") is True + assert mcp_servers_db.delete_server("srv1") is False + assert mcp_servers_db.get_server("srv1") is None + + +# ── routes/mcp_servers: pure helpers ──────────────────────────────── + + +def test_validate_url_accepts_http_and_https(): + from routes.mcp_servers import _validate_url + + assert _validate_url("http://example.com/mcp") == "http://example.com/mcp" + assert _validate_url("https://example.com/mcp") == "https://example.com/mcp" + assert _validate_url(" https://example.com/mcp ") == "https://example.com/mcp" + + +@pytest.mark.parametrize("bad", ["", " ", "ftp://x", "http://", "noscheme.com"]) +def test_validate_url_rejects_bad(bad): + from routes.mcp_servers import _validate_url + with pytest.raises(HTTPException) as exc: + _validate_url(bad) + assert exc.value.status_code == 400 + + +def test_normalize_headers(): + from routes.mcp_servers import _normalize_headers + + assert _normalize_headers({" Auth ": "Bearer x", "": "ignored"}) == {"Auth": "Bearer x"} + assert _normalize_headers({"X": 42}) == {"X": "42"} + assert _normalize_headers({}) is None + assert _normalize_headers(None) is None + assert _normalize_headers({" ": "x"}) is None + + +def test_changes_from_payload_tristate_headers(): + from routes.mcp_servers import _changes_from_payload + from models.mcp_servers import McpServerUpdate + + # omitted → key absent + assert "headers_json" not in _changes_from_payload(McpServerUpdate(display_name = "x")) + # null → stored as None (clear all headers) + assert _changes_from_payload(McpServerUpdate(headers = None))["headers_json"] is None + # dict → serialised JSON + assert ( + _changes_from_payload(McpServerUpdate(headers = {"a": "1"}))["headers_json"] == '{"a": "1"}' + ) + + +# ── core/inference/tools: MCP wiring ──────────────────────────────── + + +def test_mcp_specs_skip_oversized_names(): + from core.inference.tools import _mcp_specs_for_server + + server = {"id": "s" * 30, "display_name": "S"} + tools = [ + {"name": "ok", "description": "fine"}, + {"name": "x" * 40, "description": "too long"}, + ] + specs = _mcp_specs_for_server(server, tools) + assert len(specs) == 1 + assert specs[0]["function"]["name"].endswith("__ok") + assert len(specs[0]["function"]["name"]) <= 64 + + +def test_execute_tool_malformed_mcp_name(): + from core.inference.tools import execute_tool + out = execute_tool("mcp__no_double_underscore", {}) + assert out.startswith("Error: malformed MCP tool name") + + +def test_execute_tool_unknown_server(tmp_path, monkeypatch): + _reset_db(tmp_path, monkeypatch) + from core.inference.tools import execute_tool + assert execute_tool("mcp__missing__do_thing", {}) == "Error: MCP server 'missing' not found" + + +def test_execute_tool_disabled_server(tmp_path, monkeypatch): + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server( + id = "srv1", + display_name = "A", + url = "https://a/m", + is_enabled = False, + ) + from core.inference.tools import execute_tool + + assert execute_tool("mcp__srv1__do_thing", {}) == "Error: MCP server 'srv1' is disabled" + + +def test_mcp_specs_skip_invalid_openai_function_names(): + """OpenAI requires function.name ^[a-zA-Z0-9_-]{1,64}$; bad names 400 the request.""" + from core.inference.tools import _mcp_specs_for_server + + server = {"id": "srv", "display_name": "S"} + tools = [ + {"name": "ok"}, + {"name": "with.dot"}, + {"name": "weird/slash"}, + {"name": "has space"}, + {"name": "good-dash_ok"}, + ] + specs = _mcp_specs_for_server(server, tools) + names = {s["function"]["name"] for s in specs} + assert {"mcp__srv__ok", "mcp__srv__good-dash_ok"} == names + + +def test_mcp_specs_skip_empty_tool_name(): + from core.inference.tools import _mcp_specs_for_server + + server = {"id": "srv", "display_name": "S"} + specs = _mcp_specs_for_server(server, [{"name": "", "description": "x"}]) + assert specs == [] + + +def test_mcp_specs_drops_duplicate_names(): + """Duplicate tool names from one server -> OpenAI rejects; drop before forwarding.""" + from core.inference.tools import _mcp_specs_for_server + + server = {"id": "srv", "display_name": "S"} + tools = [{"name": "echo"}, {"name": "echo"}] + specs = _mcp_specs_for_server(server, tools) + assert len(specs) == 1 + + +def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch): + """Pre-set cancel_event -> immediate cancellation, no network round-trip.""" + import threading + from core.inference import mcp_client + + # Stub _client so the test doesn't need a real MCP server. + class _StubClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + async def call_tool(self, name, args): + import asyncio as _asyncio + await _asyncio.sleep(30) # never finishes during the test + + monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient()) + + cancel = threading.Event() + cancel.set() + out = mcp_client.call_tool_sync( + url = "https://example/mcp", + headers = None, + name = "slow", + args = {}, + timeout = 30.0, + cancel_event = cancel, + ) + assert "cancelled" in out.lower() + + +def test_clear_oauth_tokens_async_no_op_safe(tmp_path, monkeypatch): + """clear_oauth_tokens_async on a URL with no stored token must not raise; + the delete + update handlers call it best-effort regardless of state.""" + import asyncio + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_oauth_token_store", None) + asyncio.run(mcp_client.clear_oauth_tokens_async("https://example.com/mcp")) + + +def test_delete_server_calls_oauth_cleanup_when_oauth_was_on(tmp_path, monkeypatch): + """delete_mcp_server route helper must call clear_oauth_tokens_async + when the deleted row had use_oauth=true.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_oauth_token_store", None) + mcp_servers_db.create_server( + id = "oauth1", + display_name = "GH", + url = "https://gh-mcp.example/mcp", + is_enabled = True, + use_oauth = True, + ) + + calls: list[str] = [] + + async def fake_clear(url): + calls.append(url) + + monkeypatch.setattr(mcp_client, "clear_oauth_tokens_async", fake_clear) + # Patch the route's module binding too so it's seen. + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear) + asyncio.run(routes_mcp.delete_mcp_server("oauth1", current_subject = "u")) + assert calls == ["https://gh-mcp.example/mcp"] + assert mcp_servers_db.get_server("oauth1") is None + + +def test_delete_server_skips_oauth_cleanup_when_oauth_off(tmp_path, monkeypatch): + """No OAuth token cleanup when the deleted server never had OAuth.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_oauth_token_store", None) + mcp_servers_db.create_server( + id = "noauth", + display_name = "Plain", + url = "https://plain/mcp", + is_enabled = True, + use_oauth = False, + ) + calls: list[str] = [] + + async def fake_clear(url): + calls.append(url) + + monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear) + asyncio.run(routes_mcp.delete_mcp_server("noauth", current_subject = "u")) + assert calls == [] + + +def test_update_server_clears_oauth_on_url_change(tmp_path, monkeypatch): + """Changing the URL on an OAuth server must drop the old URL's tokens + so the new URL doesn't inherit credentials.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_oauth_token_store", None) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "https://old/mcp", + is_enabled = True, + use_oauth = True, + ) + calls: list[str] = [] + + async def fake_clear(url): + calls.append(url) + + monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear) + asyncio.run( + routes_mcp.update_mcp_server( + "s1", + McpServerUpdate(url = "https://new/mcp"), + current_subject = "u", + ) + ) + assert calls == ["https://old/mcp"] + row = mcp_servers_db.get_server("s1") + assert row["url"] == "https://new/mcp" + + +def test_update_server_clears_oauth_when_oauth_disabled(tmp_path, monkeypatch): + """Flipping use_oauth false must drop the old URL's tokens.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_oauth_token_store", None) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "https://u/mcp", + is_enabled = True, + use_oauth = True, + ) + calls: list[str] = [] + + async def fake_clear(url): + calls.append(url) + + monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear) + asyncio.run( + routes_mcp.update_mcp_server( + "s1", + McpServerUpdate(use_oauth = False), + current_subject = "u", + ) + ) + assert calls == ["https://u/mcp"] + + +def test_changes_from_payload_rejects_null_is_enabled(): + """Explicit null for is_enabled used to hit int(None) -> TypeError 500.""" + from routes.mcp_servers import _changes_from_payload + from models.mcp_servers import McpServerUpdate + + with pytest.raises(HTTPException) as exc: + _changes_from_payload(McpServerUpdate(is_enabled = None)) + assert exc.value.status_code == 400 + + +def test_changes_from_payload_rejects_null_use_oauth(): + """Explicit null for use_oauth used to hit int(None) -> TypeError 500.""" + from routes.mcp_servers import _changes_from_payload + from models.mcp_servers import McpServerUpdate + + with pytest.raises(HTTPException) as exc: + _changes_from_payload(McpServerUpdate(use_oauth = None)) + assert exc.value.status_code == 400 + + +def test_test_endpoint_surfaces_url_validation_as_400(tmp_path, monkeypatch): + """POST /api/mcp/servers/test must 400 on invalid URL like create/update; + it previously returned 200 with {"ok": false}.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from routes.mcp_servers import test_mcp_server + from models.mcp_servers import McpServerTestRequest + + with pytest.raises(HTTPException) as exc: + asyncio.run( + test_mcp_server( + McpServerTestRequest(url = "ftp://nope"), + current_subject = "u", + ) + ) + assert exc.value.status_code == 400 + + +def test_tool_xml_parser_handles_hyphenated_parameter_names(): + """Hyphenated property names like `issue-number` must round-trip through the + XML parser (the old `` regex dropped them).""" + from core.inference.tool_call_parser import parse_tool_calls_from_text + import json as _json + + calls = parse_tool_calls_from_text( + "" + "Bug report" + "octocat/hello" + "" + ) + assert len(calls) == 1 + args = _json.loads(calls[0]["function"]["arguments"]) + assert args == {"issue-title": "Bug report", "repo-name": "octocat/hello"} + + +def test_tool_healing_strip_handles_hyphenated_function_names(): + """core/tool_healing.py has its own copy of the XML strip regex that the + shared-parser fix missed.""" + from core.tool_healing import strip_tool_call_markup + + out = strip_tool_call_markup( + "before x after" + ) + assert out == "before after" + + +def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch): + """A tool call not in the per-request list must be refused by the GGUF + agentic loop (mirroring the safetensors path).""" + from core.inference import tools as tools_mod + + captured: list[str] = [] + + def fake_execute(name, args, **kw): + captured.append(name) + return "executed" + + monkeypatch.setattr(tools_mod, "execute_tool", fake_execute) + + # Inline allow-list check to unit-test behavior without llama-server. + def _gate(tools_advertised, called_name, args): + allowed = { + (t.get("function") or {}).get("name") + for t in (tools_advertised or []) + if (t.get("function") or {}).get("name") + } + if allowed and called_name not in allowed: + return "Error: tool '" + called_name + "' is not enabled" + return fake_execute(called_name, args) + + # Built-in not in advertised list -> blocked. + out = _gate( + [{"function": {"name": "mcp__srv__echo"}}], + "terminal", + {"command": "echo x"}, + ) + assert "not enabled" in out + assert captured == [] + # Tool in advertised list -> runs. + out = _gate( + [{"function": {"name": "mcp__srv__echo"}}], + "mcp__srv__echo", + {"text": "hi"}, + ) + assert out == "executed" + assert captured == ["mcp__srv__echo"] + + +def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch): + """Pre-set cancel_event -> no HTTP request (task used to open a transport + before the cancel check).""" + from core.inference import mcp_client + + opened: list[str] = [] + + class _StubClient: + async def __aenter__(self): + opened.append("opened") + return self + + async def __aexit__(self, *args): + return False + + async def call_tool(self, name, args): + return "ran" + + monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient()) + + import threading + + ev = threading.Event() + ev.set() + out = mcp_client.call_tool_sync( + url = "https://example/mcp", + headers = None, + name = "x", + args = {}, + timeout = 5.0, + cancel_event = ev, + ) + assert "cancelled" in out.lower() + # The client must NOT have been opened. + assert opened == [] + + +def test_clear_oauth_tokens_swallows_constructor_errors(tmp_path, monkeypatch): + """clear_oauth_tokens_async is best-effort; an OAuth constructor failure + must not bubble into a 500 from the delete/update routes.""" + import asyncio + from core.inference import mcp_client + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(mcp_client, "_oauth_token_store", None) + + # Patch the OAuth import path to raise so the entire body fails. + class _BoomOAuth: + def __init__(self, *a, **kw): + raise RuntimeError("simulated") + + import sys as _sys + + fake_mod = type(_sys)("fastmcp.client.auth") + fake_mod.OAuth = _BoomOAuth + monkeypatch.setitem(_sys.modules, "fastmcp.client.auth", fake_mod) + # Must not raise. + asyncio.run(mcp_client.clear_oauth_tokens_async("https://x/mcp")) + + +def test_tool_xml_parser_handles_hyphenated_function_names(): + """Hyphenated tool names like `mcp__srv__list-issues` must parse, else the + model can call the tool but Studio can't dispatch.""" + from core.inference.tool_call_parser import parse_tool_calls_from_text + + calls = parse_tool_calls_from_text( + "octocat/hello" + ) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "mcp__srv__list-issues" + import json as _json + + args = _json.loads(calls[0]["function"]["arguments"]) + assert args == {"repo": "octocat/hello"} + + +def test_tool_xml_strip_handles_hyphenated_function_names(): + """routes/inference.py:_TOOL_XML_RE must strip a `` + block; else hyphenated MCP tool-call XML leaks into chat history.""" + import re as _re + from pathlib import Path + + src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text() + m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL) + assert m, "could not extract _TOOL_XML_RE" + ns: dict = {"_re": _re} + exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns) + rx = ns["_TOOL_XML_RE"] + stripped = rx.sub( + "", + "before x after", + ) + assert stripped == "before after" + + +def test_safetensors_agentic_empty_allowlist_still_means_allow_all(): + """Contract: at the safetensors_agentic layer tools=[] means "no + constraint". The MCP-only-no-discovery fix lives at the route level in + inference.py, which refuses use_tools when the resolved list is empty.""" + import threading + from core.inference.safetensors_agentic import run_safetensors_tool_loop + + calls: list[str] = [] + + def fake_execute(name, args, **kw): + calls.append(name) + return "ran" + + iteration = {"n": 0} + + def fake_single_turn(messages): + iteration["n"] += 1 + if iteration["n"] == 1: + txt = '{"name":"python","arguments":{"code":"1"}}' + buf = "" + for ch in txt: + buf += ch + yield buf + else: + yield "done" + + list( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "x"}], + tools = [], + execute_tool = fake_execute, + cancel_event = threading.Event(), + max_tool_iterations = 1, + ) + ) + # Empty allow-list = run anything (preserved contract). + assert calls == [("python", {"code": "1"})] or len(calls) >= 1 diff --git a/studio/backend/tests/test_mcp_stdio_improvements.py b/studio/backend/tests/test_mcp_stdio_improvements.py new file mode 100644 index 0000000000..b0bfd45135 --- /dev/null +++ b/studio/backend/tests/test_mcp_stdio_improvements.py @@ -0,0 +1,221 @@ +"""Tests for the proposed PR #5863 improvements. + +Covers: _client() self-gating + keep_alive, OAuth normalised off for stdio +(create + update), env/header dropped on a transport-type switch, and rejecting +a command whose first token is a URL scheme. + +Run from studio/backend: python -m pytest tests/test_mcp_stdio_improvements.py -q +""" + +import asyncio + +import pytest +from fastapi import HTTPException + +from core.inference import mcp_client +from storage import mcp_servers_db + + +def _reset_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + + +def _enable(monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + + +def _disable(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False) + + +# ── P1: _client() self-gates the stdio sink ───────────────────────── + + +def test_client_refuses_stdio_when_disabled(monkeypatch): + _disable(monkeypatch) + with pytest.raises(PermissionError): + mcp_client._client("npx -y server /tmp", None) + + +def test_client_builds_stdio_when_enabled_without_spawning(monkeypatch): + _enable(monkeypatch) + # Constructing the Client must not spawn the subprocess (spawn happens on + # __aenter__); only assert it builds. + client = mcp_client._client("npx -y server /tmp", {"K": "v"}) + assert client is not None + + +def test_client_http_unaffected_by_gate(monkeypatch): + _disable(monkeypatch) + assert mcp_client._client("https://example.com/mcp", None) is not None + + +# ── P3: OAuth normalised off for stdio (create + update) ──────────── + + +def test_create_forces_oauth_off_for_stdio(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerCreate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + resp = asyncio.run( + routes_mcp.create_mcp_server( + McpServerCreate(display_name = "FS", url = "npx -y server /tmp", use_oauth = True), + current_subject = "u", + ) + ) + assert resp.use_oauth is False + assert mcp_servers_db.get_server(resp.id)["use_oauth"] == 0 + + +def test_create_keeps_oauth_for_http(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerCreate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + resp = asyncio.run( + routes_mcp.create_mcp_server( + McpServerCreate(display_name = "GH", url = "https://gh/mcp", use_oauth = True), + current_subject = "u", + ) + ) + assert resp.use_oauth is True + + +def test_update_url_to_stdio_clears_oauth(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerUpdate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + monkeypatch.setattr(mcp_client, "_oauth_token_store", None) + monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", lambda *a, **k: asyncio.sleep(0)) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp", use_oauth = True) + resp = asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "npx -y server /tmp"), current_subject = "u" + ) + ) + assert resp.use_oauth is False + + +# ── P4: env/headers dropped on a transport-type switch ────────────── + + +def test_switch_stdio_to_http_drops_env(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerUpdate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "npx server", + headers_json = '{"API_KEY": "secret"}', + ) + resp = asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "https://remote/mcp"), current_subject = "u" + ) + ) + # stdio env must NOT survive as HTTP headers on the remote endpoint + assert resp.headers == {} + assert mcp_servers_db.get_server("s1")["headers_json"] is None + + +def test_switch_keeps_explicitly_supplied_headers(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerUpdate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "npx server", + headers_json = '{"API_KEY": "secret"}', + ) + resp = asyncio.run( + routes_mcp.update_mcp_server( + "s1", + McpServerUpdate(url = "https://remote/mcp", headers = {"Authorization": "Bearer new"}), + current_subject = "u", + ) + ) + assert resp.headers == {"Authorization": "Bearer new"} + + +def test_same_transport_edit_keeps_headers(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerUpdate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "npx server", + headers_json = '{"API_KEY": "secret"}', + ) + # editing only the display name (still stdio) must keep env vars + resp = asyncio.run( + routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u") + ) + assert resp.headers == {"API_KEY": "secret"} + + +# ── P5: reject a command whose first token is a URL scheme ─────────── + + +def test_validate_url_rejects_url_scheme_command_when_enabled(monkeypatch): + from routes.mcp_servers import _validate_url + _enable(monkeypatch) + for bad in ["ftp://host/x", "file:///etc/passwd", "ws://h/y"]: + with pytest.raises(HTTPException) as exc: + _validate_url(bad) + assert exc.value.status_code == 400 + + +def test_validate_url_allows_url_in_argument(monkeypatch): + from routes.mcp_servers import _validate_url + _enable(monkeypatch) + # :// inside an ARGUMENT (not the first token) is a valid command + assert _validate_url("npx server --url https://x/mcp") == ("npx server --url https://x/mcp") + + +# ── P6: Data Recipe stdio path obeys the same host gate ───────────── +# build_mcp_providers needs the Studio-only data_designer plugin; skip if absent. + +_STDIO_RECIPE = { + "mcp_providers": [ + { + "provider_type": "stdio", + "name": "fs", + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"], + "env": {}, + } + ] +} + + +def test_data_recipe_skips_stdio_when_disabled(monkeypatch): + pytest.importorskip("data_designer") + _disable(monkeypatch) + from core.data_recipe.service import build_mcp_providers + + # gate off -> the stdio provider is dropped (no subprocess spawned) + assert build_mcp_providers(_STDIO_RECIPE) == [] + + +def test_data_recipe_builds_stdio_when_enabled(monkeypatch): + pytest.importorskip("data_designer") + _enable(monkeypatch) + from core.data_recipe.service import build_mcp_providers + + built = build_mcp_providers(_STDIO_RECIPE) + assert len(built) == 1 # constructed (not spawned) only when enabled diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py new file mode 100644 index 0000000000..c6a4898d5d --- /dev/null +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -0,0 +1,387 @@ +"""Verification tests for PR #5863 (stdio MCP server support). + +Covers the pure helpers, the route-level _validate_url gate, and that the +UNSLOTH_STUDIO_ALLOW_STDIO_MCP gate blocks the stdio transport at every +enforcement point (create/update/test/refresh/discovery/execute) when disabled +and reaches it when enabled. The transport is stubbed so no subprocess spawns; +a recorder asserts whether it was reached. +""" + +import sys + +import pytest +from fastapi import HTTPException + +from core.inference import mcp_client +from storage import mcp_servers_db + + +def _reset_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + + +def _enable(monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + + +def _disable(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False) + + +# ── transport stub + recorder ─────────────────────────────────────── + + +class _FakeTool: + def __init__(self, name): + self._name = name + + def model_dump(self, exclude_none = True): + return {"name": self._name, "description": f"{self._name} tool"} + + +class _Block: + def __init__(self, text): + self.type = "text" + self.text = text + + +class _FakeResult: + is_error = False + + def __init__(self, text): + self.content = [_Block(text)] + + +class _RecordingClient: + """Stand-in for fastmcp.Client; records that the transport was opened.""" + + def __init__(self, url, headers, use_oauth, recorder): + recorder.append({"url": url, "headers": headers, "use_oauth": use_oauth}) + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def list_tools(self): + return [_FakeTool("list_directory"), _FakeTool("write_file")] + + async def call_tool(self, name, args): + return _FakeResult(f"called {name}") + + +@pytest.fixture +def transport(monkeypatch): + """Patch mcp_client._client with a recorder. Returns the recorder list; + empty == stdio transport never reached.""" + recorder = [] + monkeypatch.setattr( + mcp_client, + "_client", + lambda url, headers, use_oauth = False: _RecordingClient(url, headers, use_oauth, recorder), + ) + return recorder + + +# ── 1. is_stdio ───────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "addr", + [ + "http://localhost:8000/mcp", + "https://example.com/mcp", + " https://example.com/mcp ", + "HTTPS://EXAMPLE.COM/mcp", + ], +) +def test_is_stdio_false_for_http(addr): + assert mcp_client.is_stdio(addr) is False + + +@pytest.mark.parametrize( + "addr", + [ + "npx -y @modelcontextprotocol/server-filesystem /tmp", + "python -m some.module", + "uvx some-server --flag", + "/usr/local/bin/my-server", + ], +) +def test_is_stdio_true_for_commands(addr): + assert mcp_client.is_stdio(addr) is True + + +# ── 2. parse_stdio_command ────────────────────────────────────────── + + +def test_parse_basic_argv(): + assert mcp_client.parse_stdio_command( + "npx -y @modelcontextprotocol/server-filesystem /tmp" + ) == ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp"] + + +def test_parse_keeps_url_argument_as_one_command(): + # gemini "high": a :// inside an ARGUMENT must not break the command. + assert mcp_client.parse_stdio_command("npx server --endpoint https://example.com/mcp") == [ + "npx", + "server", + "--endpoint", + "https://example.com/mcp", + ] + + +def test_parse_quoted_arg(): + assert mcp_client.parse_stdio_command('python -m mod --name "a b"') == [ + "python", + "-m", + "mod", + "--name", + "a b", + ] + + +def test_parse_empty_returns_empty_list(): + assert mcp_client.parse_stdio_command(" ") == [] + + +def test_parse_unclosed_quote_raises_valueerror(): + with pytest.raises(ValueError): + mcp_client.parse_stdio_command('npx "unclosed') + + +def test_parse_windows_strips_wrapping_quotes(monkeypatch): + # gemini "medium": posix=False keeps backslash paths but also the + # wrapping quotes; the PR strips a matched pair so argv[0] is clean. + monkeypatch.setattr(sys, "platform", "win32") + parts = mcp_client.parse_stdio_command(r'"C:\Program Files\node\node.exe" server.js') + assert parts[0] == r"C:\Program Files\node\node.exe" + assert parts[1] == "server.js" + + +# ── 3. stdio_mcp_enabled ──────────────────────────────────────────── + + +@pytest.mark.parametrize("val", ["0", "false", "true", "", " 1 ", "yes", "2"]) +def test_stdio_disabled_for_non_exact_one(monkeypatch, val): + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", val) + assert mcp_client.stdio_mcp_enabled() is False + + +def test_stdio_enabled_only_for_exact_one(monkeypatch): + _disable(monkeypatch) + assert mcp_client.stdio_mcp_enabled() is False + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + assert mcp_client.stdio_mcp_enabled() is True + + +# ── 4. probe_timeout ──────────────────────────────────────────────── + + +def test_probe_timeout_matrix(): + assert mcp_client.probe_timeout("https://x/mcp", False) == 8.0 + assert mcp_client.probe_timeout("https://x/mcp", True) == 305.0 + assert mcp_client.probe_timeout("npx server", False) == 60.0 + # oauth wins regardless of address kind (documented behaviour) + assert mcp_client.probe_timeout("npx server", True) == 305.0 + + +# ── 5. _validate_url gate ─────────────────────────────────────────── + + +def test_validate_url_gate_off_rejects_stdio(monkeypatch): + _disable(monkeypatch) + from routes.mcp_servers import _validate_url + + assert _validate_url("https://example.com/mcp") == "https://example.com/mcp" + # urlparse reads "localhost:8000" scheme as "localhost", so it lands here too. + for bad in [ + "npx server", + "python -m mod", + "ftp://host", + "example.com", + "localhost:8000", + r"C:\node\node.exe server.js", + ]: + with pytest.raises(HTTPException) as exc: + _validate_url(bad) + assert exc.value.status_code == 400 + + +def test_validate_url_gate_off_message_depends_on_whitespace(monkeypatch): + # The message names a command only when the value has whitespace, and + # never says "desktop app only" (self-hosted can opt in via the env var). + _disable(monkeypatch) + from routes.mcp_servers import _validate_url + + with pytest.raises(HTTPException) as exc: + _validate_url("npx -y @modelcontextprotocol/server-filesystem /tmp") + cmd = exc.value.detail.lower() + assert "http://" in cmd and "https://" in cmd + assert "local command" in cmd + assert "desktop app" not in cmd + + with pytest.raises(HTTPException) as exc: + _validate_url("example.com") + lone = exc.value.detail.lower() + assert "http://" in lone and "https://" in lone + assert "local command" not in lone + + +def test_validate_url_gate_on_accepts_stdio(monkeypatch): + _enable(monkeypatch) + from routes.mcp_servers import _validate_url + + assert _validate_url("npx -y server /tmp") == "npx -y server /tmp" + # http still works when stdio is on + assert _validate_url("https://x/mcp") == "https://x/mcp" + # url-bearing argument accepted as a command + assert _validate_url("npx server --url https://x/mcp") == ("npx server --url https://x/mcp") + # A lone token is ambiguous; accept it as a command rather than + # guessing it's a URL (no regression for single binaries). + assert _validate_url("/usr/local/bin/my-mcp-server") == "/usr/local/bin/my-mcp-server" + assert _validate_url("mcp-server-sqlite") == "mcp-server-sqlite" + # empty / unparseable still rejected + for bad in [" ", '"unclosed']: + with pytest.raises(HTTPException) as exc: + _validate_url(bad) + assert exc.value.status_code == 400 + + +# ── 6. gate enforcement at every spawn path (mocked transport) ────── + + +def test_create_route_gate(tmp_path, monkeypatch, transport): + import asyncio + + from models.mcp_servers import McpServerCreate + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + payload = McpServerCreate(display_name = "FS", url = "npx -y server /tmp") + + _disable(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(routes_mcp.create_mcp_server(payload, current_subject = "u")) + assert exc.value.status_code == 400 + + _enable(monkeypatch) + resp = asyncio.run(routes_mcp.create_mcp_server(payload, current_subject = "u")) + assert resp.url == "npx -y server /tmp" + + +def test_update_http_to_stdio_blocked_when_off(tmp_path, monkeypatch): + import asyncio + + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + _disable(monkeypatch) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp") + # editing url -> stdio command must 400 (http->stdio bypass closed) + with pytest.raises(HTTPException) as exc: + asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "npx server"), current_subject = "u" + ) + ) + assert exc.value.status_code == 400 + + +def test_test_route_gate(tmp_path, monkeypatch, transport): + import asyncio + + from models.mcp_servers import McpServerTestRequest + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + req = McpServerTestRequest(url = "npx -y server /tmp") + + _disable(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(routes_mcp.test_mcp_server(req, current_subject = "u")) + assert exc.value.status_code == 400 + assert transport == [] # transport never opened + + _enable(monkeypatch) + res = asyncio.run(routes_mcp.test_mcp_server(req, current_subject = "u")) + assert res.ok and res.tool_count == 2 + assert len(transport) == 1 + + +def test_refresh_route_gate(tmp_path, monkeypatch, transport): + import asyncio + + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + # a stdio row, as if carried over from a desktop DB + mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server") + + _disable(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u")) + assert exc.value.status_code == 400 + assert transport == [] + + _enable(monkeypatch) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u")) + assert res.ok and res.tool_count == 2 + assert len(transport) == 1 + + +def test_discovery_gate(tmp_path, monkeypatch, transport): + import asyncio + + from core.inference.tools import get_enabled_mcp_tools + + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True) + + _disable(monkeypatch) + assert asyncio.run(get_enabled_mcp_tools()) == [] + assert transport == [] # filtered out before any probe + + _enable(monkeypatch) + specs = asyncio.run(get_enabled_mcp_tools()) + assert len(specs) == 2 + assert len(transport) == 1 + + +def test_execute_gate(tmp_path, monkeypatch, transport): + from core.inference.tools import execute_tool + + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True) + + _disable(monkeypatch) + out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"}) + assert "disabled on this host" in out + assert transport == [] + + _enable(monkeypatch) + out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"}) + assert out == "called list_directory" + assert len(transport) == 1 + + +# ── 7. env vars ride headers_json as the subprocess env ───────────── + + +def test_stdio_env_passed_through(tmp_path, monkeypatch, transport): + from core.inference.tools import execute_tool + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + mcp_servers_db.create_server( + id = "stdio1", + display_name = "FS", + url = "npx server", + headers_json = '{"API_KEY": "sk-test"}', + is_enabled = True, + ) + execute_tool("mcp__stdio1__list_directory", {}) + assert transport[-1]["headers"] == {"API_KEY": "sk-test"} diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index bbaf20298d..1005431926 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -24,21 +24,25 @@ if str(_BACKEND_ROOT) not in sys.path: @pytest.fixture(scope = "module") def main_module(): import main as _main # noqa: F401 - return _main -# ===================================================================== # MaxBodyMiddleware -# ===================================================================== -def _make_protected_app(max_bytes: int, main_module): +def _make_protected_app( + max_bytes: int, + main_module, + upload_passthrough_prefixes: tuple = (), + upload_passthrough_max_bytes_getter = None, +): app = FastAPI() app.add_middleware( main_module.MaxBodyMiddleware, - max_bytes = max_bytes, - protected_prefixes = ("/v1/chat/completions", "/api/train"), + max_bytes_getter = lambda: max_bytes, + protected_prefixes = ("/v1/chat/completions", "/api/settings", "/api/train"), + upload_passthrough_prefixes = upload_passthrough_prefixes, + upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter, ) @app.post("/v1/chat/completions") @@ -49,6 +53,20 @@ def _make_protected_app(max_bytes: int, main_module): async def other(payload: dict): return {"ok": True, "unprotected": True} + @app.put("/api/settings/upload-limit") + async def update_upload_limit(payload: dict): + return {"ok": True, "limit": payload.get("max_upload_size_mb")} + + @app.post("/api/train/upload") + async def upload(request: Request): + total = 0 + chunks = 0 + async for chunk in request.stream(): + if chunk: + chunks += 1 + total += len(chunk) + return {"ok": True, "chunks": chunks, "total": total} + @app.get("/api/train/status") async def status_get(): return {"ok": True, "get": True} @@ -78,9 +96,19 @@ class TestMaxBodyMiddleware: assert r.status_code == 200 assert r.json()["unprotected"] is True + def test_settings_put_body_over_cap_rejected(self, main_module): + app = _make_protected_app(1024, main_module) + c = TestClient(app) + r = c.put( + "/api/settings/upload-limit", + json = {"max_upload_size_mb": 500, "padding": "x" * 5000}, + ) + assert r.status_code == 413 + assert "too large" in r.json()["detail"].lower() + def test_chunked_upload_over_cap_rejected(self, main_module): - # Regression: declared-Content-Length-only check could be bypassed - # by chunked transfer-encoding. + # Regression: declared-Content-Length-only check could be bypassed by + # chunked transfer-encoding. app = _make_protected_app(1024, main_module) c = TestClient(app) @@ -121,10 +149,61 @@ class TestMaxBodyMiddleware: r = c.get("/api/train/status") assert r.status_code == 200 + def test_upload_passthrough_uses_dedicated_declared_cap(self, main_module): + app = _make_protected_app( + 128, + main_module, + upload_passthrough_prefixes = ("/api/train/upload",), + upload_passthrough_max_bytes_getter = lambda: 1024, + ) + c = TestClient(app) + r = c.post( + "/api/train/upload", + content = b"x" * 512, + headers = {"content-type": "application/octet-stream"}, + ) + assert r.status_code == 200 + assert r.json()["total"] == 512 + + def test_upload_passthrough_rejects_declared_body_over_dedicated_cap(self, main_module): + app = _make_protected_app( + 128, + main_module, + upload_passthrough_prefixes = ("/api/train/upload",), + upload_passthrough_max_bytes_getter = lambda: 256, + ) + c = TestClient(app) + r = c.post( + "/api/train/upload", + content = b"x" * 512, + headers = {"content-type": "application/octet-stream"}, + ) + assert r.status_code == 413 + assert "256" in r.json()["detail"] + + def test_upload_passthrough_requires_content_length(self, main_module): + app = _make_protected_app( + 128, + main_module, + upload_passthrough_prefixes = ("/api/train/upload",), + upload_passthrough_max_bytes_getter = lambda: 1024, + ) + c = TestClient(app) + + def gen(): + yield b"x" * 64 + yield b"y" * 64 + + r = c.post( + "/api/train/upload", + content = gen(), + headers = {"content-type": "application/octet-stream"}, + ) + assert r.status_code == 411 + assert "Content-Length" in r.json()["detail"] + -# ===================================================================== # SecurityHeadersMiddleware / CSP -# ===================================================================== def _make_csp_app(main_module, attach_nonce: str | None = None): @@ -174,7 +253,10 @@ class TestSecurityHeadersMiddleware: assert r.headers["x-frame-options"] == "DENY" assert r.headers["x-content-type-options"] == "nosniff" assert r.headers["referrer-policy"] == "no-referrer" - assert "camera=()" in r.headers["permissions-policy"] + permissions_policy = r.headers["permissions-policy"] + assert "camera=()" in permissions_policy + assert "microphone=(self)" in permissions_policy + assert "geolocation=()" in permissions_policy assert r.headers["server"] == "unsloth-studio" def test_internal_nonce_header_is_spliced_into_csp_and_stripped(self, main_module): @@ -185,9 +267,7 @@ class TestSecurityHeadersMiddleware: csp = r.headers["content-security-policy"] assert f"'nonce-{nonce}'" in csp # Internal handoff header must not leak to clients. - assert main_module._CSP_SCRIPT_NONCE_HEADER not in { - k.lower() for k in r.headers.keys() - } + assert main_module._CSP_SCRIPT_NONCE_HEADER not in {k.lower() for k in r.headers.keys()} def test_build_csp_helper_shape(self, main_module): plain = main_module._build_csp() @@ -196,32 +276,25 @@ class TestSecurityHeadersMiddleware: nonced = main_module._build_csp("XYZ") assert "script-src 'self' 'nonce-XYZ';" in nonced - def test_img_src_allows_google_favicons(self, main_module): - # sources.tsx fetches https://www.google.com/s2/favicons?... ; without - # this allowlist entry citation favicons fall back to gray initials. + def test_img_and_media_allow_https_sources(self, main_module): + # Model-card READMEs and citation favicons pull images/media from many + # https origins (HF LFS/XET CDNs, shields/badge hosts, GitHub-hosted + # assets, audio/video samples). img-src/media-src allow any https source + # so they render; this mirrors the desktop CSP in tauri.conf.json. csp = main_module._build_csp() - img_directive = next( - chunk.strip() + directives = { + chunk.strip().split()[0]: chunk.strip().split() for chunk in csp.split(";") - if chunk.strip().startswith("img-src ") - ) - # Tokenise and compare with `==` so CodeQL's URL-substring rule does - # not read directive-string `in` membership as URL sanitisation. - img_sources = img_directive.split() - assert any(src == "https://www.google.com" for src in img_sources) - # Pre-existing favicon CDNs stay allowed. - for host in ( - "https://t0.gstatic.com", - "https://t1.gstatic.com", - "https://t2.gstatic.com", - "https://t3.gstatic.com", - ): - assert any(src == host for src in img_sources) + if chunk.strip() + } + for name in ("img-src", "media-src"): + assert name in directives, f"missing {name} directive" + # Tokenise and compare with `==` so CodeQL's URL-substring rule does + # not read directive-string `in` membership as URL sanitisation. + assert any(src == "https:" for src in directives[name]) -# ===================================================================== # /api/health auth gate -# ===================================================================== @pytest.fixture @@ -250,9 +323,9 @@ def health_app(tmp_path, monkeypatch): class TestHealthAuthGate: - # Launcher / frontend bootstrap fields are available unauth so the Tauri - # watchdog can re-adopt a sibling backend and the SPA can detect chat-only - # mode before any token exists. Version / device_type still require a bearer. + # Launcher / frontend bootstrap fields are unauth so the Tauri watchdog can + # re-adopt a sibling backend and the SPA can detect chat-only mode before + # any token exists. Version / device_type still require a bearer. LAUNCHER_BITS = ( "service", "studio_root_id", @@ -279,7 +352,7 @@ class TestHealthAuthGate: assert forbidden not in body def test_invalid_bearer_returns_launcher_bits_only(self, health_app): - # Regression: calling the async dep without await made any Bearer header pass. + # Regression: calling the async dep without await let any Bearer header pass. c = TestClient(health_app) r = c.get( "/api/health", diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 16cca7dd40..9871965ce8 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -103,8 +103,7 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite( - monkeypatch, - tmp_path, + monkeypatch, tmp_path ): _install_fake_mlx(monkeypatch) calls = [] @@ -160,10 +159,9 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri assert isinstance(backend._tokenizer, _DummyTokenizer) -# Regression: MLXInferenceBackend.generate_chat_response must accept the -# four template kwargs (tools / enable_thinking / reasoning_effort / -# preserve_thinking) so the route layer can forward what the user -# toggled in the UI. The previous signature raised +# Regression: generate_chat_response must accept the four template kwargs +# (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route +# layer can forward UI toggles. The old signature raised # "got an unexpected keyword argument 'tools'" on Mac. @@ -185,8 +183,8 @@ def test_mlx_generate_chat_response_accepts_template_kwargs(): def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): - """The Mac text path must route through apply_chat_template_for_ - generation so reasoning / tool kwargs reach the tokenizer.""" + """Mac text path must route through apply_chat_template_for_generation so + reasoning / tool kwargs reach the tokenizer.""" _install_fake_mlx(monkeypatch) from core.inference.mlx_inference import MLXInferenceBackend @@ -199,14 +197,13 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): return "" monkeypatch.setattr( - "core.inference.chat_template_helpers." "apply_chat_template_for_generation", + "core.inference.chat_template_helpers.apply_chat_template_for_generation", _fake_apply, raising = True, ) - # mlx_lm.stream_generate yields response objects with .token; make a - # one-token generator so _generate_text returns without touching the - # real stack. + # mlx_lm.stream_generate yields response objects with .token; use a + # one-token generator so _generate_text returns without the real stack. import types as _types mlx_lm_pkg = _types.ModuleType("mlx_lm") @@ -228,7 +225,11 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): class _Tok: chat_template = "x" - def decode(self, ids, skip_special_tokens = False): + def decode( + self, + ids, + skip_special_tokens = False, + ): return "hi" backend = MLXInferenceBackend() @@ -247,7 +248,7 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): ) ) assert out == ["hi"] - # The kwargs the user toggled must reach the chat-template helper. + # The toggled kwargs must reach the chat-template helper. assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}] assert captured["kwargs"]["enable_thinking"] is True assert captured["kwargs"]["reasoning_effort"] == "medium" diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 98c7bdaa55..4402031467 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -45,12 +45,8 @@ def _load_worker_module(): setattr(wheel_utils, name, lambda *_args, **_kwargs: None) sys.modules["utils.wheel_utils"] = wheel_utils - worker_path = ( - Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" - ) - spec = importlib.util.spec_from_file_location( - "mlx_training_worker_under_test", worker_path - ) + worker_path = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" + spec = importlib.util.spec_from_file_location("mlx_training_worker_under_test", worker_path) module = importlib.util.module_from_spec(spec) assert spec.loader is not None spec.loader.exec_module(module) @@ -66,6 +62,11 @@ def _load_worker_module(): _worker = _load_worker_module() _normalize_mlx_studio_optimizer = _worker._normalize_mlx_studio_optimizer _normalize_mlx_studio_scheduler = _worker._normalize_mlx_studio_scheduler +_mlx_vlm_max_resized_size = _worker._mlx_vlm_max_resized_size +_mlx_vlm_resized_image_layout = _worker._mlx_vlm_resized_image_layout +_copy_mlx_vlm_image_processor = _worker._copy_mlx_vlm_image_processor +_resize_mlx_vlm_image = _worker._resize_mlx_vlm_image +_adapt_for_mlx_vlm = _worker._adapt_for_mlx_vlm def test_mlx_studio_optimizer_aliases_are_explicit(): @@ -82,3 +83,134 @@ def test_mlx_studio_rejects_unknown_optimizer(): def test_mlx_studio_rejects_unknown_scheduler(): with pytest.raises(ValueError, match = "Unsupported LR scheduler for MLX training"): _normalize_mlx_studio_scheduler("linear_typo") + + +def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer(): + assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256) + assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512) + assert _mlx_vlm_max_resized_size(1000, 1000, 512) == (512, 512) + assert _mlx_vlm_max_resized_size(256, 128, 1536) == (256, 128) + assert _mlx_vlm_max_resized_size(512, 256, 512) == (512, 256) + # Half-pixel cases must match the Torch collator (not banker's round). + assert _mlx_vlm_max_resized_size(333, 1000, 500) == (167, 500) + assert _mlx_vlm_max_resized_size(1000, 333, 500) == (500, 167) + + +def test_mlx_vlm_resize_keeps_default_numpy_layout_hwc(): + Image = pytest.importorskip("PIL.Image") + image = Image.new("RGB", (320, 200), color = (10, 20, 30)) + + resized = _resize_mlx_vlm_image(image, 128) + + assert resized.shape == (80, 128, 3) + assert resized.flags.c_contiguous + + +def test_mlx_vlm_resize_uses_requested_chw_numpy_layout(): + Image = pytest.importorskip("PIL.Image") + image = Image.new("RGB", (320, 200), color = (10, 20, 30)) + + resized = _resize_mlx_vlm_image(image, 128, image_layout = "chw") + + assert resized.shape == (3, 80, 128) + assert resized.flags.c_contiguous + + +def test_mlx_vlm_resized_image_layout_probes_processor_contract(): + class ChwOnlyImageProcessor: + def __call__(self, images = None): + image = images[0] + if image.shape[0] == 3: + return {"pixel_values": image} + raise ValueError("expected CHW") + + class HwcImageProcessor: + def __call__(self, images = None): + image = images[0] + if image.shape[-1] == 3: + return {"pixel_values": image} + raise ValueError("expected HWC") + + assert ( + _mlx_vlm_resized_image_layout( + types.SimpleNamespace(image_processor = ChwOnlyImageProcessor()) + ) + == "chw" + ) + assert ( + _mlx_vlm_resized_image_layout(types.SimpleNamespace(image_processor = HwcImageProcessor())) + is None + ) + + +def test_mlx_vlm_layout_probe_copies_image_processor(): + class StatefulImageProcessor: + def __init__(self): + self.calls = 0 + + def __call__(self, images = None): + self.calls += 1 + image = images[0] + if image.shape[0] == 3: + return {"pixel_values": image} + raise ValueError("expected CHW") + + image_processor = StatefulImageProcessor() + + layout = _mlx_vlm_resized_image_layout(types.SimpleNamespace(image_processor = image_processor)) + + assert layout == "chw" + assert image_processor.calls == 0 + + +def test_mlx_vlm_image_processor_copy_refuses_uncopyable_processors(): + class UncopyableImageProcessor: + def __copy__(self): + raise RuntimeError("no copy") + + def __deepcopy__(self, _memo): + raise RuntimeError("no deepcopy") + + image_processor = UncopyableImageProcessor() + + assert _copy_mlx_vlm_image_processor(image_processor) is None + + +def test_mlx_vlm_layout_probe_skips_uncopyable_processors(): + class UncopyableImageProcessor: + def __copy__(self): + raise RuntimeError("no copy") + + def __deepcopy__(self, _memo): + raise RuntimeError("no deepcopy") + + def __call__(self, images = None): + raise AssertionError("live processor should not be probed") + + assert ( + _mlx_vlm_resized_image_layout( + types.SimpleNamespace(image_processor = UncopyableImageProcessor()) + ) + is None + ) + + +def test_mlx_vlm_adapter_applies_chw_layout_to_message_images(): + Image = pytest.importorskip("PIL.Image") + image = Image.new("RGB", (320, 200), color = (10, 20, 30)) + item = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": "Describe it."}, + ], + } + ] + } + + adapted = _adapt_for_mlx_vlm([item], resize = 128, image_layout = "chw") + + assert adapted[0]["image"].shape == (3, 80, 128) + assert adapted[0]["messages"][0]["content"][0] == {"type": "image"} diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py index 3481e29948..417ec74d17 100644 --- a/studio/backend/tests/test_models_get_model_config_case_resolution.py +++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py @@ -50,9 +50,7 @@ def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch): return _DummyModelConfig() monkeypatch.setattr(models_route, "is_local_path", lambda _: False) - monkeypatch.setattr( - models_route, "resolve_cached_repo_id_case", lambda _: "Org/Model" - ) + monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda _: "Org/Model") monkeypatch.setattr(models_route, "load_model_defaults", _record_load) monkeypatch.setattr(models_route, "is_vision_model", _record_vision) monkeypatch.setattr(models_route, "is_embedding_model", _record_embedding) diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py new file mode 100644 index 0000000000..d5e8d13652 --- /dev/null +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Separate-file MTP drafter (Gemma 4) contracts. + +Pins: the drafter-path predicate and its two layering mirrors, Gemma +effective-size extraction, companion classification in variant plans +(including resume from pre-fix manifests where the drafter leaked into a +quant's main files), and local drafter detection / self-pairing rejection. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from hub.utils.download_manifest import ExpectedFile +from hub.utils.gguf import is_mtp_drafter_path +from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files +from utils.models.model_config import ( + _is_mtp_drafter, + detect_gguf_model, + detect_mtp_file, + extract_model_size_b, +) + + +# ── Predicate + layering mirrors ───────────────────────────────────── + +DRAFTER_CASES = [ + ("mtp-gemma-4-12b-it.gguf", True), + ("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True), + ("foo/MTP/bar.gguf", True), + ("gemma-4-12b-it-Q8_0.gguf", False), + # Baked-in Qwen MTP repos: the head is inside the main GGUF, the file + # IS the model -- must never be classified as a companion. + ("Qwen3.6-27B-MTP-Q4_K_M.gguf", False), + ("prompt-mtp-test.gguf", False), + ("smtp/model.gguf", False), + ("mtp-readme.txt", False), +] + + +@pytest.mark.parametrize("path,expected", DRAFTER_CASES) +def test_drafter_predicate_and_mirrors_agree(path, expected): + from core.inference.llama_cpp import _is_companion_gguf_path + + assert is_mtp_drafter_path(path) is expected + assert _is_mtp_drafter(path) is expected + # The core mirror bundles mmproj; none of these inputs are mmproj, so + # it must agree with the canonical predicate. + assert _is_companion_gguf_path(path) is expected + + +# ── Gemma effective-size extraction ────────────────────────────────── + + +@pytest.mark.parametrize( + "model_id,size_b", + [ + ("unsloth/gemma-4-E2B-it-GGUF", 2.0), + ("unsloth/gemma-4-E4B-it", 4.0), + ("unsloth/gemma-3n-E4B-it", 4.0), + # MoE active params beat effective and total notation. + ("unsloth/Qwen3.5-35B-A3B", 3.0), + ("unsloth/gemma-4-12b-it-GGUF", 12.0), + ("unsloth/Qwen3.5-9B-MTP-GGUF", 9.0), + ("no-size-here", None), + ], +) +def test_extract_model_size_b(model_id, size_b): + assert extract_model_size_b(model_id) == size_b + + +# ── Variant plan companion classification ──────────────────────────── + + +def _sib(name: str, size: int, sha: str): + return SimpleNamespace(rfilename = name, size = size, lfs = {"sha256": sha}) + + +GEMMA_SIBLINGS = [ + _sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"), + _sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"), + _sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"), + _sib("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", 100, "mtp-sub-q8"), + _sib("MTP/gemma-4-12b-it-BF16-MTP.gguf", 200, "mtp-sub-bf16"), + _sib("mmproj-F16.gguf", 500, "mmproj"), +] + + +def test_variant_plans_carry_drafter_as_companion(): + plans = build_gguf_variant_plans(GEMMA_SIBLINGS) + + # No phantom quants from the drafter's Q8_0 label or the MTP/ copies. + assert set(plans) == {"q4_k_m", "q8_0"} + for plan in plans.values(): + assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames + assert not any("MTP/" in name for name in plan.target_filenames) + assert "drafter" in plan.companion_hashes + assert "drafter" not in plan.main_hashes + assert plan.mmproj_filenames == frozenset({"mmproj-F16.gguf"}) + + q4 = plans["q4_k_m"] + assert q4.main_filenames == frozenset({"gemma-4-12b-it-Q4_K_M.gguf"}) + assert q4.main_size_bytes == 4_000 + # Download size = main + mmproj + drafter. + assert q4.download_size_bytes == 4_600 + + +def test_baked_in_repo_plans_unchanged(): + plans = build_gguf_variant_plans([_sib("Qwen3.6-27B-MTP-Q4_K_M.gguf", 4_000, "q4")]) + assert plans["q4_k_m"].target_filenames == ("Qwen3.6-27B-MTP-Q4_K_M.gguf",) + + +def test_old_manifest_resume_reclassifies_drafter(): + # Pre-fix manifests could leak the drafter into a quant's expected + # files; resume must classify it as a companion, not a main shard. + old = [ + ExpectedFile(path = "gemma-4-12b-it-Q8_0.gguf", size = 8_000, sha256 = "main-q8"), + ExpectedFile(path = "mtp-gemma-4-12b-it.gguf", size = 100, sha256 = "drafter"), + ] + plan = plan_from_expected_files("Q8_0", old) + assert plan.main_hashes == frozenset({"main-q8"}) + assert plan.companion_hashes == frozenset({"drafter"}) + assert plan.mmproj_filenames == frozenset() + + +# ── Local detection / self-pairing ─────────────────────────────────── + + +def test_detect_mtp_file_finds_root_sibling(tmp_path): + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"x") + (tmp_path / "mtp-model.gguf").write_bytes(b"x") + (tmp_path / "MTP").mkdir() + (tmp_path / "MTP" / "model-Q8_0-MTP.gguf").write_bytes(b"x") + + found = detect_mtp_file(str(tmp_path / "model-Q4_K_M.gguf")) + assert found is not None + assert found.endswith("mtp-model.gguf") + + +def test_detect_mtp_file_none_without_sibling(tmp_path): + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"x") + assert detect_mtp_file(str(tmp_path / "model-Q4_K_M.gguf")) is None + + +def test_detect_gguf_model_rejects_drafter_file(tmp_path): + drafter = tmp_path / "mtp-model.gguf" + drafter.write_bytes(b"x") + assert detect_gguf_model(str(drafter)) is None + + +def test_detect_gguf_model_dir_skips_companions(tmp_path): + main = tmp_path / "model-Q4_K_M.gguf" + main.write_bytes(b"xxxx") + # Companions are larger so a size-sorted pick would wrongly win. + (tmp_path / "mtp-model.gguf").write_bytes(b"x" * 64) + (tmp_path / "mmproj-F16.gguf").write_bytes(b"x" * 128) + + assert detect_gguf_model(str(tmp_path)) == str(main.resolve()) + + +def test_detect_mtp_file_pairs_by_weight_name(tmp_path): + # Multi-model folder: each weight must get its own drafter, never the + # first-sorted foreign one. + (tmp_path / "gemma-4-12b-it-Q4_K_M.gguf").write_bytes(b"x") + (tmp_path / "gemma-4-31B-it-Q4_K_M.gguf").write_bytes(b"x") + (tmp_path / "mtp-gemma-4-12b-it.gguf").write_bytes(b"x") + (tmp_path / "mtp-gemma-4-31B-it.gguf").write_bytes(b"x") + + found = detect_mtp_file(str(tmp_path / "gemma-4-31B-it-Q4_K_M.gguf")) + assert found is not None and found.endswith("mtp-gemma-4-31B-it.gguf") + + +def test_detect_mtp_file_skips_foreign_drafter(tmp_path): + (tmp_path / "qwen3-8b-Q4_K_M.gguf").write_bytes(b"x") + (tmp_path / "mtp-gemma-4-12b-it.gguf").write_bytes(b"x") + assert detect_mtp_file(str(tmp_path / "qwen3-8b-Q4_K_M.gguf")) is None + + +def test_detect_mtp_file_qat_prefix_layout(tmp_path): + # unsloth's qat repo: drafter stem omits the -qat suffix but prefixes + # the weight name (mtp-gemma-4-12B-it.gguf / gemma-4-12B-it-qat-Q4_0.gguf). + (tmp_path / "gemma-4-12B-it-qat-Q4_0.gguf").write_bytes(b"x") + (tmp_path / "mtp-gemma-4-12B-it.gguf").write_bytes(b"x") + found = detect_mtp_file(str(tmp_path / "gemma-4-12B-it-qat-Q4_0.gguf")) + assert found is not None and found.endswith("mtp-gemma-4-12B-it.gguf") + + +def test_detect_mtp_file_search_root(tmp_path): + # Weight in a quant subdir, drafter at the granted directory root. + sub = tmp_path / "Q4_K_M" + sub.mkdir() + (sub / "gemma-4-12b-it-Q4_K_M.gguf").write_bytes(b"x") + (tmp_path / "mtp-gemma-4-12b-it.gguf").write_bytes(b"x") + found = detect_mtp_file(str(sub / "gemma-4-12b-it-Q4_K_M.gguf"), search_root = str(tmp_path)) + assert found is not None and found.endswith("mtp-gemma-4-12b-it.gguf") + + +# ── Reload dedup includes the drafter ──────────────────────────────── + + +def _loaded_backend(weight, drafter_path): + from core.inference.llama_cpp import LlamaCppBackend + + b = LlamaCppBackend() + # Shape matches atexit cleanup expectations (terminate/wait/kill). + b._process = SimpleNamespace( + poll = lambda: None, + terminate = lambda: None, + wait = lambda timeout = None: 0, + kill = lambda: None, + ) + b._healthy = True + b._model_identifier = "local-gemma" + b._gguf_path = str(weight) + b._hf_variant = None + b._requested_n_ctx = 4096 + b._cache_type_kv = None + b._requested_spec_mode = "auto" + b._speculative_type = "draft-mtp" if drafter_path else "default" + b._spec_draft_n_max = None + b._chat_template_override = None + b._extra_args = None + b._mtp_draft_path = drafter_path + return b + + +def _target_state_kwargs(weight, mtp_draft_path): + return dict( + model_identifier = "local-gemma", + hf_variant = None, + n_ctx = 4096, + cache_type_kv = None, + speculative_type = "auto", + spec_draft_n_max = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + gguf_path = str(weight), + mtp_draft_path = mtp_draft_path, + ) + + +def test_already_in_target_state_bounces_on_new_drafter(tmp_path): + weight = tmp_path / "gemma-4-12b-it-Q4_K_M.gguf" + weight.write_bytes(b"x") + drafter = tmp_path / "mtp-gemma-4-12b-it.gguf" + drafter.write_bytes(b"x") + + # Loaded without a drafter; one now exists on disk -> must reload. + b = _loaded_backend(weight, None) + assert not b._already_in_target_state(**_target_state_kwargs(weight, str(drafter))) + # Same drafter as launched -> still deduped. + b = _loaded_backend(weight, str(drafter)) + assert b._already_in_target_state(**_target_state_kwargs(weight, str(drafter))) + + +def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path): + # Direct selection of an MTP/ copy: the basename alone has no mtp- + # prefix, so rejection relies on the parent dir name. + sub = tmp_path / "MTP" + sub.mkdir() + copy = sub / "gemma-4-12b-it-BF16-MTP.gguf" + copy.write_bytes(b"x") + assert detect_gguf_model(str(copy)) is None + # Selecting the MTP dir itself must not surface the copies as models. + assert detect_gguf_model(str(sub)) is None diff --git a/studio/backend/tests/test_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py index 4d7528d238..5cd7c876cc 100644 --- a/studio/backend/tests/test_multimodal_document.py +++ b/studio/backend/tests/test_multimodal_document.py @@ -3,17 +3,16 @@ """Tests for PDF / document attachment translation on external providers. -Studio introduces a normalised `input_document` content part on -ChatCompletionRequest so the frontend doesn't have to know the -per-provider attachment shape: +Studio adds a normalised `input_document` content part on +ChatCompletionRequest so the frontend needn't know the per-provider +attachment shape: -- Anthropic: translates to `{type:"document", source:{type:"base64"|"url", ...}}` -- OpenAI Responses: translates to `{type:"input_file", file_data|file_url, filename?}` +- Anthropic: `{type:"document", source:{type:"base64"|"url", ...}}` +- OpenAI Responses: `{type:"input_file", file_data|file_url, filename?}` -These tests pin the translation shape on both paths for base64 data -URIs and remote URLs, with optional filename metadata, and confirm -unknown / empty document parts are dropped without breaking the -request. +Pins the translation shape on both paths for base64 data URIs and remote +URLs (with optional filename), and confirms unknown / empty document +parts are dropped without breaking the request. """ import asyncio @@ -86,10 +85,8 @@ _PDF_DATA_URI = f"data:application/pdf;base64,{_TINY_PDF_B64}" def _strip_cache(p: dict) -> dict: - # Studio's prompt-cache wiring attaches cache_control:{type:ephemeral} - # to the tail block of the last user message; strip it before - # comparing the document core fields so this test stays focused - # on the translation, not the caching layer. + # Strip the prompt-cache cache_control off the last user block so this + # test focuses on translation, not the caching layer. return {k: v for k, v in p.items() if k != "cache_control"} @@ -117,8 +114,8 @@ def test_anthropic_base64_pdf_becomes_document_block(monkeypatch): types = [p.get("type") for p in parts] assert "document" in types, parts doc = _strip_cache(next(p for p in parts if p.get("type") == "document")) - # citations: {enabled: true} opts into Anthropic's natural-citation - # pipeline; without it the citations_delta handler is a no-op. + # citations:{enabled:true} opts into Anthropic's citation pipeline; + # without it the citations_delta handler is a no-op. assert doc == { "type": "document", "source": { @@ -179,9 +176,8 @@ def test_anthropic_empty_document_part_is_dropped(monkeypatch): def test_anthropic_empty_only_document_drops_whole_message(monkeypatch): - # If the ONLY part in a user message is an unparseable input_document, - # the helper must NOT append an empty-content message to the outbound - # body (Anthropic 400s on "at least one block is required"). + # If the only part is an unparseable input_document, the helper must not + # append an empty-content message (Anthropic 400s on "at least one block"). captured = _capture( monkeypatch, provider = "anthropic", @@ -192,14 +188,13 @@ def test_anthropic_empty_only_document_drops_whole_message(monkeypatch): ], ) msgs = captured["body"]["messages"] - # The empty-content message must be skipped; only the second remains. + # Empty-content message skipped; only the second remains. assert len(msgs) == 1, msgs def test_anthropic_empty_data_uri_payload_is_dropped(monkeypatch): - # Codex P2: `data:application/pdf;base64,` with no payload (or - # whitespace-only) would create an empty `source.data` that - # Anthropic 400s on. Must be filtered before the wire. + # A `data:application/pdf;base64,` with empty/whitespace payload makes an + # empty `source.data` that Anthropic 400s on; filter it before the wire. captured = _capture( monkeypatch, provider = "anthropic", @@ -228,12 +223,9 @@ def test_anthropic_empty_data_uri_payload_is_dropped(monkeypatch): def test_anthropic_empty_data_uri_falls_back_to_file_url(monkeypatch): - # Codex P2 follow-up: my previous fix added the empty-data-URI -> - # file_url fallback to the OpenAI side but missed the Anthropic - # side, where the empty-payload branch did `continue` and discarded - # an otherwise-valid file_url on the same part. Mirror the OpenAI - # behavior so a malformed inline payload + remote URL still - # attaches. + # The empty-data-URI -> file_url fallback existed on OpenAI but not + # Anthropic, which discarded a valid file_url on the same part. Mirror + # OpenAI so a malformed inline payload + remote URL still attaches. captured = _capture( monkeypatch, provider = "anthropic", @@ -255,7 +247,7 @@ def test_anthropic_empty_data_uri_falls_back_to_file_url(monkeypatch): ) parts = captured["body"]["messages"][0]["content"] doc = _strip_cache(next(p for p in parts if p.get("type") == "document")) - # base64 source MUST NOT have landed on the wire; URL source survived. + # base64 source MUST NOT reach the wire; URL source survives. assert doc == { "type": "document", "source": {"type": "url", "url": "https://example.com/doc.pdf"}, @@ -317,11 +309,7 @@ def test_openai_base64_pdf_becomes_input_file(monkeypatch): user_msg = captured["body"]["input"][0] parts = user_msg["content"] fileblk = next(p for p in parts if p.get("type") == "input_file") - assert fileblk == { - "type": "input_file", - "file_data": _PDF_DATA_URI, - "filename": "paper.pdf", - } + assert fileblk == {"type": "input_file", "file_data": _PDF_DATA_URI, "filename": "paper.pdf"} def test_openai_url_pdf_becomes_input_file(monkeypatch): @@ -344,18 +332,13 @@ def test_openai_url_pdf_becomes_input_file(monkeypatch): ) parts = captured["body"]["input"][0]["content"] fileblk = next(p for p in parts if p.get("type") == "input_file") - assert fileblk == { - "type": "input_file", - "file_url": "https://example.com/doc.pdf", - } + assert fileblk == {"type": "input_file", "file_url": "https://example.com/doc.pdf"} def test_openai_empty_data_uri_falls_back_to_file_url(monkeypatch): - # Codex P2 follow-up: an empty `data:application/pdf;base64,` - # payload was being preferred over a perfectly valid `file_url` - # in the same part, sending `file_data=""` to OpenAI and 400ing - # the whole turn. The translator must treat empty data URIs as - # missing and recover via file_url. + # An empty `data:application/pdf;base64,` payload was preferred over a valid + # `file_url` in the same part, sending `file_data=""` and 400ing. The + # translator must treat empty data URIs as missing and recover via file_url. captured = _capture( monkeypatch, provider = "openai", @@ -377,7 +360,7 @@ def test_openai_empty_data_uri_falls_back_to_file_url(monkeypatch): ) parts = captured["body"]["input"][0]["content"] fileblk = next(p for p in parts if p.get("type") == "input_file") - # file_data MUST NOT be on the wire; file_url survives. + # file_data MUST NOT reach the wire; file_url survives. assert "file_data" not in fileblk, fileblk assert fileblk["file_url"] == "https://example.com/doc.pdf" assert fileblk["filename"] == "doc.pdf" @@ -409,8 +392,8 @@ def test_openai_whitespace_only_data_uri_falls_back_to_file_url(monkeypatch): def test_openai_empty_data_uri_without_fallback_is_dropped(monkeypatch): - # If the only signal is an empty data URI (no file_url), the - # whole part is skipped rather than sent as `file_data=""`. + # Only signal is an empty data URI (no file_url): skip the whole part + # rather than send `file_data=""`. captured = _capture( monkeypatch, provider = "openai", @@ -455,13 +438,9 @@ def test_openai_empty_document_part_is_dropped(monkeypatch): # ── Pydantic schema + builder pass-through ────────────────────────── -# -# The translation tests above call the external-provider client directly -# with hand-built dicts, which bypasses BOTH ChatCompletionRequest's -# discriminated Union AND routes/inference._build_external_messages. The -# tests below close that gap: parse an input_document part through the -# real request schema, run the builder, and assert the part survives to -# the dict the client would receive. +# The tests above call the client with hand-built dicts, bypassing the schema +# and _build_external_messages. The tests below parse an input_document part +# through the real schema + builder and assert it survives to the client dict. def test_chat_message_accepts_input_document_part(): @@ -489,10 +468,9 @@ def test_chat_message_accepts_input_document_part(): def test_build_external_messages_passes_input_document_for_anthropic_and_openai(): - # Both providers' stream helpers have explicit input_document - # translation logic (Anthropic -> {type:"document"}, OpenAI - # Responses -> {type:"input_file"}), so the part round-trips - # through the builder unchanged on those routes. + # Both providers' stream helpers translate input_document (Anthropic -> + # {type:"document"}, OpenAI Responses -> {type:"input_file"}), so the + # part round-trips through the builder unchanged on those routes. from models.inference import ChatMessage from routes.inference import _build_external_messages @@ -512,9 +490,7 @@ def test_build_external_messages_passes_input_document_for_anthropic_and_openai( ) ] for provider in ("anthropic", "openai"): - out = _build_external_messages( - msgs, supports_vision = True, provider_type = provider - ) + out = _build_external_messages(msgs, supports_vision = True, provider_type = provider) assert len(out) == 1, (provider, out) parts = out[0]["content"] assert parts[0] == {"type": "text", "text": "summarise"}, provider @@ -527,10 +503,10 @@ def test_build_external_messages_passes_input_document_for_anthropic_and_openai( def test_build_external_messages_strips_input_document_for_unmapped_providers(): # Codex P1 follow-up: gemini / mistral / kimi / openrouter / deepseek - # / custom go through generic /chat/completions passthrough that - # forwards `messages` verbatim. Handing them an `input_document` - # part fails the upstream validator. Builder must strip the part - # for every provider whose stream helper doesn't translate it. + # / custom use generic /chat/completions passthrough that forwards + # `messages` verbatim, so an `input_document` part fails the upstream + # validator. The builder must strip it for any provider whose stream + # helper doesn't translate it. from models.inference import ChatMessage from routes.inference import _build_external_messages @@ -550,9 +526,7 @@ def test_build_external_messages_strips_input_document_for_unmapped_providers(): ) ] for provider in ("gemini", "mistral", "kimi", "openrouter", "deepseek", "qwen"): - out = _build_external_messages( - msgs, supports_vision = True, provider_type = provider - ) + out = _build_external_messages(msgs, supports_vision = True, provider_type = provider) assert len(out) == 1, (provider, out) parts = out[0]["content"] types = [p.get("type") for p in parts if isinstance(p, dict)] @@ -562,8 +536,8 @@ def test_build_external_messages_strips_input_document_for_unmapped_providers(): def test_build_external_messages_strips_input_document_when_provider_type_unknown(): - # Defensive: legacy callers that don't pass provider_type must - # not leak the part to an unknown destination. + # Defensive: legacy callers without provider_type must not leak the + # part to an unknown destination. from models.inference import ChatMessage from routes.inference import _build_external_messages diff --git a/studio/backend/tests/test_native_context_length.py b/studio/backend/tests/test_native_context_length.py index 60622c776d..de1ca0649e 100644 --- a/studio/backend/tests/test_native_context_length.py +++ b/studio/backend/tests/test_native_context_length.py @@ -3,11 +3,11 @@ """Tests for the native_context_length feature (PR #4746). -Verifies that the new `native_context_length` property on LlamaCppBackend -and the corresponding Pydantic model fields work correctly. The raw GGUF -`_context_length` must never be overwritten by VRAM-capping logic. +Verifies the `native_context_length` property on LlamaCppBackend and the +matching Pydantic fields. The raw GGUF `_context_length` must never be +overwritten by VRAM-capping logic. -Requires no GPU, network, or external libraries beyond pytest and pydantic. +Needs no GPU, network, or libraries beyond pytest and pydantic. """ import io @@ -21,8 +21,8 @@ from unittest.mock import patch import pytest # --------------------------------------------------------------------------- -# Stub heavy / unavailable external dependencies before importing the -# module under test. Same pattern as test_kv_cache_estimation.py. +# Stub heavy / unavailable deps before importing the module under test. +# Same pattern as test_kv_cache_estimation.py. # --------------------------------------------------------------------------- _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) @@ -38,7 +38,7 @@ sys.modules.setdefault("loggers", _loggers_stub) _structlog_stub = _types.ModuleType("structlog") sys.modules.setdefault("structlog", _structlog_stub) -# httpx -- stub only the names referenced at import / class-definition time +# httpx -- stub only names referenced at import / class-definition time _httpx_stub = _types.ModuleType("httpx") for _exc_name in ( "ConnectError", @@ -227,7 +227,7 @@ class TestContextValueSeparation: def test_all_equal_when_uncapped(self, backend): """All three equal when no VRAM constraint.""" backend._context_length = 8192 - # No effective or max set -- properties fall back to _context_length + # No effective/max set -- properties fall back to _context_length. assert backend.native_context_length == 8192 assert backend.max_context_length == 8192 assert backend.context_length == 8192 @@ -241,16 +241,16 @@ class TestContextValueSeparation: backend._embedding_length = 4096 original = backend._context_length - # Simulate a very small VRAM budget that forces capping + # Tiny VRAM budget forces capping. result = backend._fit_context_to_vram( requested_ctx = 131072, available_mib = 512, # very small model_size_bytes = 0, ) - # _fit_context_to_vram returns the capped value, not modifying _context_length + # Returns the capped value without modifying _context_length. assert backend._context_length == original assert backend.native_context_length == original - # The returned capped value should be <= requested + # Capped value must be <= requested. assert result <= 131072 def test_native_gt_context_when_capped(self, backend): @@ -271,6 +271,7 @@ class TestPydanticModels: def test_load_response_has_field(self): """Field exists in LoadResponse.model_fields.""" assert "native_context_length" in LoadResponse.model_fields + assert "context_length" in LoadResponse.model_fields def test_load_response_defaults_none(self): """Omitting native_context_length defaults to None.""" @@ -319,6 +320,7 @@ class TestPydanticModels: def test_status_response_has_field(self): """Field exists in InferenceStatusResponse.model_fields.""" assert "native_context_length" in InferenceStatusResponse.model_fields + assert "context_length" in InferenceStatusResponse.model_fields def test_status_response_has_chat_template_field(self): """Status includes chat_template so the UI can rehydrate after refresh.""" @@ -332,9 +334,7 @@ class TestPydanticModels: def test_status_response_chat_template_roundtrip(self): """chat_template serializes and validates as part of status.""" resp = InferenceStatusResponse(chat_template = "{{ messages }}") - roundtripped = InferenceStatusResponse.model_validate_json( - resp.model_dump_json() - ) + roundtripped = InferenceStatusResponse.model_validate_json(resp.model_dump_json()) assert roundtripped.chat_template == "{{ messages }}" def test_roundtrip_preserves_value(self): @@ -349,6 +349,18 @@ class TestPydanticModels: roundtripped = LoadResponse.model_validate_json(resp.model_dump_json()) assert roundtripped.native_context_length == 131072 + def test_context_length_roundtrip(self): + """Runtime context_length serializes for non-GGUF/hub models.""" + resp = LoadResponse( + status = "loaded", + model = "test", + display_name = "Test", + inference = {}, + context_length = 8192, + ) + roundtripped = LoadResponse.model_validate_json(resp.model_dump_json()) + assert roundtripped.context_length == 8192 + # ===================================================================== # D. TestRouteCompleteness -- source-level verification @@ -372,7 +384,7 @@ class TestRouteCompleteness: start = self._source.find(f"{class_name}(", idx) if start == -1: break - # Find matching closing paren (simple depth counter) + # Find the matching closing paren via a depth counter. depth = 0 end = start for i, ch in enumerate(self._source[start:], start): @@ -390,9 +402,7 @@ class TestRouteCompleteness: def test_gguf_load_responses_have_field(self): """Every GGUF LoadResponse (is_gguf = True) includes native_context_length.""" blocks = self._find_construction_blocks("LoadResponse") - gguf_blocks = [ - b for b in blocks if "is_gguf = True" in b or "is_gguf=True" in b - ] + gguf_blocks = [b for b in blocks if "is_gguf = True" in b or "is_gguf=True" in b] assert ( len(gguf_blocks) >= 2 ), f"Expected at least 2 GGUF LoadResponse blocks, found {len(gguf_blocks)}" @@ -404,16 +414,24 @@ class TestRouteCompleteness: def test_non_gguf_load_responses_omit_field(self): """Non-GGUF LoadResponse blocks do not set native_context_length (defaults to None).""" blocks = self._find_construction_blocks("LoadResponse") - non_gguf = [ - b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b - ] - # Non-GGUF paths should not reference native_context_length - # (Pydantic defaults it to None, so not setting it is correct) + non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b] + # Non-GGUF paths shouldn't reference native_context_length + # (Pydantic defaults it to None, so omitting it is correct). for block in non_gguf: assert ( "native_context_length" not in block ), f"Non-GGUF LoadResponse should not set native_context_length:\n{block[:200]}" + def test_non_gguf_load_responses_set_runtime_context_length(self): + """Non-GGUF LoadResponse blocks report runtime context_length.""" + blocks = self._find_construction_blocks("LoadResponse") + non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b] + assert non_gguf, "Expected at least one non-GGUF LoadResponse block" + for block in non_gguf: + assert ( + "context_length" in block + ), f"Non-GGUF LoadResponse should set context_length:\n{block[:200]}" + def test_status_path(self): """InferenceStatusResponse construction with llama_backend has the field.""" blocks = self._find_construction_blocks("InferenceStatusResponse") @@ -422,7 +440,24 @@ class TestRouteCompleteness: if "llama_backend" in block and "native_context_length" in block: found = True break - assert found, "No InferenceStatusResponse block with llama_backend has native_context_length" + assert ( + found + ), "No InferenceStatusResponse block with llama_backend has native_context_length" + + def test_non_gguf_status_path_reports_runtime_context_length(self): + """Non-GGUF InferenceStatusResponse reports context_length from model_info.""" + blocks = self._find_construction_blocks("InferenceStatusResponse") + found = False + for block in blocks: + if "is_gguf = False" in block and "context_length" in block: + found = True + break + assert found, "No non-GGUF InferenceStatusResponse block with context_length" + + def test_openai_models_listing_reports_context_length(self): + """/v1/models includes context_length when the backend knows it.""" + assert 'entry["context_length"]' in self._source + assert 'model_info.get("context_length")' in self._source # ===================================================================== @@ -480,7 +515,7 @@ class TestNativeContextEdgeCases: backend._read_gguf_metadata(path) assert backend.native_context_length == 131072 - # Simulate VRAM capping by setting effective and max + # Simulate VRAM capping via effective and max. backend._effective_context_length = 16384 backend._max_context_length = 32768 assert backend.native_context_length == 131072 diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index d3b2f553a2..aab58adfff 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -3,29 +3,14 @@ """Regression tests for the offline GGUF cache fallback path (#5505). -Three failure modes hit users when ``huggingface.co`` is unreachable -but the requested GGUF repo is fully cached locally: +When ``huggingface.co`` is unreachable but the repo is cached, three failures +hit: ``list_gguf_variants`` 500'd (empty dropdown), ``detect_gguf_model_remote`` +returned None (GGUF-only repo misrouted), and ``_download_gguf`` synthesised a +name absent from cache. Follow-ups: the cache filter matches the snapshot-relative +path (subdir layouts findable), and DNS auto-detect scopes ``HF_HUB_OFFLINE`` to +one load so a transient hiccup can't pin the singleton offline. -* ``list_gguf_variants`` raised through ``HTTPException(500)`` so the - variant dropdown sat empty. -* ``detect_gguf_model_remote`` returned ``None`` so a GGUF-only repo - was misrouted into the transformers/Unsloth backend (on macOS this - surfaced as a hardware error). -* ``_download_gguf`` fell back to a synthetic ``{repo}-{variant}.gguf`` - name that did not exist in cache when the in-repo filename did not - echo the repo name (e.g. ``unsloth/Qwen3.6-27B-MTP-GGUF`` ships - ``Qwen3.6-27B-UD-Q4_K_XL.gguf`` with no ``MTP`` token). - -Two follow-up regressions covered here: - -* P1 #1: the cache-side variant filter must match the snapshot-relative - path, not just the basename, so subdir layouts like - ``BF16/foo.gguf`` are findable. -* P1 #2: the DNS auto-detect must scope ``HF_HUB_OFFLINE`` to one load - via try/finally so a transient resolver hiccup cannot lock the - long-lived ``LlamaCppBackend`` singleton offline forever. - -No GPU, no network, no subprocess. Linux, macOS, Windows compatible. +No GPU, no network, no subprocess. Linux/macOS/Windows compatible. """ from __future__ import annotations @@ -44,8 +29,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Stub heavy/unavailable external deps before importing the modules -# under test (same pattern as other studio backend tests). +# Stub heavy/unavailable external deps before importing the modules under +# test (same pattern as other studio backend tests). _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) @@ -149,8 +134,7 @@ def _siblings(items: dict[str, int]): """Mock ``hf_model_info(...).siblings`` payload.""" return _types.SimpleNamespace( siblings = [ - _types.SimpleNamespace(rfilename = name, size = size) - for name, size in items.items() + _types.SimpleNamespace(rfilename = name, size = size) for name, size in items.items() ], ) @@ -174,12 +158,8 @@ class TestIterHfCacheSnapshots: assert list(_iter_hf_cache_snapshots("unsloth/bare")) == [] def test_yields_newest_first(self, hf_cache): - old = _build_cache( - hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40 - ) - new = _build_cache( - hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40 - ) + old = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) + new = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) os.utime(old, (1000, 1000)) os.utime(new, (2000, 2000)) out = list(_iter_hf_cache_snapshots("unsloth/multi")) @@ -187,7 +167,7 @@ class TestIterHfCacheSnapshots: def test_repo_id_match_is_case_insensitive(self, hf_cache): _build_cache(hf_cache, "unsloth/Foo-GGUF", {"Foo-Q4_K_M.gguf": 1}) - # Lookup with a different casing of the org/name still resolves + # Lookup with different org/name casing still resolves out = list(_iter_hf_cache_snapshots("UNSLOTH/foo-gguf")) assert len(out) == 1 @@ -218,9 +198,7 @@ class TestListGgufVariantsFromCache: class TestListGgufVariantsOffline: - def test_offline_env_short_circuits_api( - self, hf_cache, clean_offline_env, monkeypatch - ): + def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1}) monkeypatch.setenv("HF_HUB_OFFLINE", "1") @@ -232,11 +210,7 @@ class TestListGgufVariantsOffline: assert len(variants) == 1 assert variants[0].quant == "UD-Q4_K_XL" - def test_api_exception_falls_back_to_cache( - self, - hf_cache, - clean_offline_env, - ): + def test_api_exception_falls_back_to_cache(self, hf_cache, clean_offline_env): _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1}) def boom(*a, **k): @@ -282,10 +256,8 @@ class TestDetectGgufFromCache: assert _detect_gguf_from_hf_cache("unsloth/a") == "a-UD-Q4_K_XL.gguf" def test_subdir_only_quant_resolves(self, hf_cache): - """P1 #1 regression: ``BF16/foo.gguf`` (quant only in directory). - Before the fix, the offline cache scan matched on basename and - missed this layout, falling through to the synthetic - ``{repo}-{variant}.gguf`` heuristic.""" + """Regression: ``BF16/foo.gguf`` (quant only in directory). The pre-fix + cache scan matched on basename and missed this layout.""" _build_cache( hf_cache, "unsloth/gpt-oss-20b-BF16", @@ -302,12 +274,7 @@ class TestDetectGgufFromCache: class TestDetectGgufModelRemoteOffline: - def test_offline_env_short_circuits_retries( - self, - hf_cache, - clean_offline_env, - monkeypatch, - ): + def test_offline_env_short_circuits_retries(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1}) monkeypatch.setenv("HF_HUB_OFFLINE", "1") @@ -331,12 +298,8 @@ class TestDetectGgufModelRemoteOffline: out = detect_gguf_model_remote("unsloth/a") assert out == "a-Q4_K_M.gguf" - def test_repository_not_found_does_not_consult_cache( - self, - hf_cache, - clean_offline_env, - ): - # Cache has a file but the API explicitly says repo is gone. + def test_repository_not_found_does_not_consult_cache(self, hf_cache, clean_offline_env): + # Cache has a file but the API says the repo is gone. _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1}) class RepositoryNotFoundError(Exception): @@ -430,12 +393,7 @@ class TestHfOfflineIfDnsDead: assert did_set is False assert "HF_HUB_OFFLINE" not in os.environ - def test_user_set_hf_hub_offline_is_preserved( - self, - dns, - clean_offline_env, - monkeypatch, - ): + def test_user_set_hf_hub_offline_is_preserved(self, dns, clean_offline_env, monkeypatch): # User explicitly set offline before launching Studio. monkeypatch.setenv("HF_HUB_OFFLINE", "1") dns.fail() @@ -445,12 +403,7 @@ class TestHfOfflineIfDnsDead: # Helper must not pop a variable it did not set. assert os.environ.get("HF_HUB_OFFLINE") == "1" - def test_user_set_transformers_offline_is_preserved( - self, - dns, - clean_offline_env, - monkeypatch, - ): + def test_user_set_transformers_offline_is_preserved(self, dns, clean_offline_env, monkeypatch): monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") dns.fail() with _hf_offline_if_dns_dead(): @@ -461,11 +414,7 @@ class TestHfOfflineIfDnsDead: # TRANSFORMERS_OFFLINE pre-existed -> preserved. assert os.environ.get("TRANSFORMERS_OFFLINE") == "1" - def test_exception_inside_block_still_restores_env( - self, - dns, - clean_offline_env, - ): + def test_exception_inside_block_still_restores_env(self, dns, clean_offline_env): dns.fail() with pytest.raises(RuntimeError, match = "boom"): with _hf_offline_if_dns_dead(): @@ -476,9 +425,8 @@ class TestHfOfflineIfDnsDead: class TestExtractQuantLabelSubdir: - """``_extract_quant_label`` must consider the parent directories when - the basename has no quant token. Subdir layouts like ``BF16/foo.gguf`` - are documented in this codebase and surface through the cache scan.""" + """``_extract_quant_label`` must consider parent dirs when the basename has + no quant token (subdir layouts like ``BF16/foo.gguf``).""" def test_quant_in_basename_unchanged(self): assert _extract_quant_label("BF16/foo-BF16.gguf") == "BF16" @@ -491,22 +439,15 @@ class TestExtractQuantLabelSubdir: assert _extract_quant_label("UD-Q4_K_XL/weight.gguf") == "UD-Q4_K_XL" def test_deeper_nesting_picks_nearest_quant_dir(self): - # When multiple parent segments could match, prefer the one closest - # to the file (innermost). This matches how repos like - # ``models/MXFP4_MOE/foo.gguf`` are laid out. + # Multiple matching parents: prefer the innermost (closest to the file). assert _extract_quant_label("models/MXFP4_MOE/foo.gguf") == "MXFP4_MOE" class TestDownloadMmprojOfflineCacheFallback: - """``LlamaCppBackend._download_mmproj`` must resolve cached mmproj - GGUFs offline, same shape as ``_download_gguf``. Without this the - offline vision GGUF load path returns ``None`` even when the mmproj - is present in cache.""" + """``_download_mmproj`` must resolve cached mmproj GGUFs offline, like + ``_download_gguf``; else the offline vision load returns None despite a cache hit.""" - def test_cache_lookup_returns_cached_mmproj_when_list_repo_files_fails( - self, - hf_cache, - ): + def test_cache_lookup_returns_cached_mmproj_when_list_repo_files_fails(self, hf_cache): _build_cache( hf_cache, "unsloth/vision-GGUF", @@ -520,7 +461,12 @@ class TestDownloadMmprojOfflineCacheFallback: def boom_list(*a, **k): raise OSError("offline") - def fake_download(*, repo_id, filename, token = None): + def fake_download( + *, + repo_id, + filename, + token = None, + ): # Echo back so the test can verify the cache-resolved filename return f"/fake/cache/{repo_id}/{filename}" @@ -551,7 +497,12 @@ class TestDownloadMmprojOfflineCacheFallback: captured = {} - def fake_download(*, repo_id, filename, token = None): + def fake_download( + *, + repo_id, + filename, + token = None, + ): captured["filename"] = filename return f"/fake/{filename}" @@ -586,7 +537,7 @@ class TestDownloadMmprojOfflineCacheFallback: class TestListLocalGgufVariantsSubdir: """Subdir layouts like ``BF16/foo.gguf`` and ``Q4_K_M/foo.gguf`` must - produce distinct quant labels, not collapse on basename.""" + yield distinct quant labels, not collapse on basename.""" def test_two_subdir_variants_do_not_collapse(self, tmp_path): from utils.models.model_config import list_local_gguf_variants @@ -655,9 +606,7 @@ class TestListGgufVariantsPermanentErrors: list_gguf_variants("u/gated-gguf") assert type(exc_info.value).__name__ == "GatedRepoError" - def test_transient_error_still_falls_back_to_cache( - self, hf_cache, clean_offline_env - ): + def test_transient_error_still_falls_back_to_cache(self, hf_cache, clean_offline_env): from utils.models.model_config import list_gguf_variants _build_cache(hf_cache, "u/transient-gguf", {"foo-Q4_K_M.gguf": 1}) @@ -671,12 +620,11 @@ class TestListGgufVariantsPermanentErrors: class TestDetectGgufFromCacheExcludesMmproj: - """A partial cache with only a vision projector must not route the - projector as the main model.""" + """A partial cache with only a vision projector must not route it as + the main model.""" def test_mmproj_only_returns_none(self, hf_cache): from utils.models.model_config import _detect_gguf_from_hf_cache - _build_cache( hf_cache, "u/vision-only-mmproj", @@ -701,9 +649,8 @@ class TestDetectGgufFromCacheExcludesMmproj: class TestProbeDnsDeadNoGlobalTimeoutMutation: - """``_probe_dns_dead`` must not change ``socket.setdefaulttimeout`` - process-wide -- concurrent sockets without explicit timeout would - inherit it for the probe window.""" + """``_probe_dns_dead`` must not change ``socket.setdefaulttimeout`` process-wide; + concurrent sockets would inherit it during the probe window.""" def test_default_timeout_unchanged_when_dns_up(self, monkeypatch): import socket as _socket @@ -724,7 +671,7 @@ class TestProbeDnsDeadNoGlobalTimeoutMutation: try: _probe_dns_dead("example.invalid", timeout = 0.5) finally: - # Restore exact state regardless of any test-side mutation. + # Restore exact state regardless of test-side mutation. original_set(prev) assert set_calls == [], ( @@ -739,7 +686,6 @@ class TestProbeDnsDeadNoGlobalTimeoutMutation: # Simulate a wedged resolver: thread blocks forever. def wedged(host): import threading - threading.Event().wait() monkeypatch.setattr(_socket, "gethostbyname", wedged) @@ -747,9 +693,8 @@ class TestProbeDnsDeadNoGlobalTimeoutMutation: class TestWaitForHealthRetriesOnReadError: - """A TCP RST mid-read while llama-server is still binding the port - (Windows: WinError 10054) must not abort the health-poll loop -- - that masks a legitimate 'still warming up' state as a fatal load.""" + """A TCP RST mid-read while llama-server is still binding (Windows: WinError + 10054) must not abort the health-poll loop and mask warmup as a fatal load.""" def test_read_error_then_success(self, monkeypatch): import httpx diff --git a/studio/backend/tests/test_offline_inference_parent.py b/studio/backend/tests/test_offline_inference_parent.py index 088be4fcd5..71331220d6 100644 --- a/studio/backend/tests/test_offline_inference_parent.py +++ b/studio/backend/tests/test_offline_inference_parent.py @@ -103,10 +103,7 @@ class TestEnvOffline: class TestTransformersVersionOfflineShortCircuits: def test_tokenizer_config_skips_urllib_when_offline( - self, - monkeypatch, - clean_offline_env, - tmp_path, + self, monkeypatch, clean_offline_env, tmp_path ): # No local config + offline env -> must NOT call urlopen. monkeypatch.setenv("HF_HUB_OFFLINE", "1") @@ -118,12 +115,7 @@ class TestTransformersVersionOfflineShortCircuits: with patch("urllib.request.urlopen", boom): assert _check_tokenizer_config_needs_v5(unique) is False - def test_config_550_skips_urllib_when_offline( - self, - monkeypatch, - clean_offline_env, - tmp_path, - ): + def test_config_550_skips_urllib_when_offline(self, monkeypatch, clean_offline_env, tmp_path): monkeypatch.setenv("HF_HUB_OFFLINE", "1") unique = f"unsloth/never-cached-{tmp_path.name}-cfg" @@ -139,9 +131,7 @@ class TestLoraDetectOffline: OfflineModeIsEnabled; cached adapter_config.json wins.""" def test_hf_model_info_short_circuits_with_OfflineModeIsEnabled( - self, - monkeypatch, - clean_offline_env, + self, monkeypatch, clean_offline_env ): from unittest.mock import MagicMock @@ -150,7 +140,7 @@ class TestLoraDetectOffline: monkeypatch.setenv("HF_HUB_OFFLINE", "1") # Studio catches Exception broadly; pin that the call still happens - # (so cached LoRAs aren't missed) and returns fast via mock. + # (so cached LoRAs aren't missed) and returns fast via the mock. class _OfflineModeIsEnabled(Exception): pass @@ -171,10 +161,7 @@ class TestLoraDetectOffline: ) def test_cached_lora_detected_when_api_unreachable( - self, - monkeypatch, - clean_offline_env, - tmp_path, + self, monkeypatch, clean_offline_env, tmp_path ): """A cached adapter_config.json must still mark the repo as a LoRA when the HF API is unreachable.""" diff --git a/studio/backend/tests/test_openai_citation_markers.py b/studio/backend/tests/test_openai_citation_markers.py index ccc17be329..3549a5993e 100644 --- a/studio/backend/tests/test_openai_citation_markers.py +++ b/studio/backend/tests/test_openai_citation_markers.py @@ -4,8 +4,8 @@ """Tests for the OpenAI Responses-API citation marker rewriter. The stream interleaves text deltas with ``\\ue200cite\\ue202SOURCE_ID\\ue201`` -markers. The rewriter resolves each to `[N](URL)` when the annotation has -arrived and drops it otherwise; the URL list still flows to Sources via +markers. The rewriter resolves each to `[N](URL)` once the annotation arrives +and drops it otherwise; the URL list still flows to Sources via `_record_url_citation`. Reference: https://developers.openai.com/api/docs/guides/citation-formatting @@ -58,8 +58,7 @@ def test_marker_rewritten_to_link_when_annotation_known(): def test_unknown_source_marker_dropped_silently(): text = f"Foo {_marker('turn9view9')} bar." out = _replace_openai_citation_markers(text, []) - # Marker stripped, no garbled "E202" glyph leaks through, and the - # surrounding text stays intact. + # Marker stripped, no garbled "E202" glyph leaks, surrounding text intact. assert not _has_marker_codepoints(out) assert "E202" not in out assert "turn9view9" not in out @@ -67,8 +66,7 @@ def test_unknown_source_marker_dropped_silently(): def test_multiple_concatenated_markers_resolved_in_order(): - """Real-world wire shape: a string of markers butted up against each other - after a sentence, as in the user-reported bug.""" + """Real-world wire shape: markers butted together after a sentence (user-reported bug).""" markers = "".join(_marker(f"turn{i}view{j}") for i, j in [(1, 0), (1, 1), (3, 0)]) text = f"All animals ranked. {markers}" citations = [ @@ -136,8 +134,7 @@ def test_citation_without_source_id_does_not_crash(citation): def test_multiple_source_id_aliases_resolve_to_same_url(): - """Every alias for the same URL must resolve, not just the first. - Regression for the Codex P1 on the original PR.""" + """Every alias for the same URL must resolve, not just the first (Codex P1 regression).""" a = _marker("turn0view0") b = _marker("turn0view0_span_1") c = _marker("turn0view0_span_2") @@ -150,15 +147,14 @@ def test_multiple_source_id_aliases_resolve_to_same_url(): }, ] out = _replace_openai_citation_markers(text, citations) - # All three aliases collapse onto citation [1] -- the URL is the - # same so it would be misleading to show three different numbers. + # All three aliases collapse onto citation [1] -- same URL, so showing + # three different numbers would mislead. assert out.count("[[1]](https://example.com/paris)") == 3 assert not _has_marker_codepoints(out) def test_source_ids_list_and_legacy_source_id_both_resolve(): - """Mixed-shape citation: legacy ``source_id`` plus newer - ``source_ids`` aliases both resolve.""" + """Mixed-shape citation: legacy ``source_id`` plus newer ``source_ids`` aliases both resolve.""" legacy = _marker("legacy_id") alias = _marker("alias_id") text = f"Both {legacy} and {alias} work." @@ -174,11 +170,9 @@ def test_source_ids_list_and_legacy_source_id_both_resolve(): assert not _has_marker_codepoints(out) -# --------------------------------------------------------------------------- # _rewrite_citation_markers_partial: deferred-annotation tests. OpenAI emits -# url_citation annotations on a subsequent SSE event; this helper reports +# url_citation annotations on a later SSE event; this helper reports # `has_unresolved` so the stream loop defers emission. See PR #5713 audit. -# --------------------------------------------------------------------------- def test_partial_known_marker_resolves_and_clears_unresolved(): @@ -202,7 +196,7 @@ def test_partial_unknown_marker_preserves_verbatim_and_flags(): def test_partial_resolves_after_late_annotation(): - """Two-pass: first call sees no citations, second resolves after annotation.""" + """Two-pass: first call sees no citations; second resolves after annotation.""" text = f"See {_marker('s1')} for details." out1, unresolved1 = _rewrite_citation_markers_partial(text, []) assert unresolved1 is True @@ -214,17 +208,15 @@ def test_partial_resolves_after_late_annotation(): def test_partial_multi_source_partial_resolution_keeps_marker_pending(): - """Any unresolved token in a multi-source marker leaves the whole marker - verbatim with ``unresolved`` True; defer until every id resolves or - end-of-stream forces a flush (dropping unresolved tokens then).""" + """Any unresolved token in a multi-source marker leaves the whole marker verbatim with ``unresolved`` True until every id resolves or end-of-stream flushes.""" cite = f"{CITE_START}cite{CITE_DELIM}known{CITE_DELIM}locator{CITE_STOP}" text = f"Pre {cite} post." citations = [{"source_id": "known", "url": "https://example.com/y"}] out, unresolved = _rewrite_citation_markers_partial(text, citations) assert unresolved is True assert cite in out - # End-of-stream force flush: drop the unresolved token, keep the - # resolved link. The streamer routes pending segments through + # End-of-stream force flush: drop the unresolved token, keep the resolved + # link. The streamer routes pending segments through # `_replace_openai_citation_markers` at force=True for this. forced = _replace_openai_citation_markers(out, citations) assert "[[1]](https://example.com/y)" in forced diff --git a/studio/backend/tests/test_openai_citation_markers_edge.py b/studio/backend/tests/test_openai_citation_markers_edge.py index ffe8c6b6eb..f44975ce33 100644 --- a/studio/backend/tests/test_openai_citation_markers_edge.py +++ b/studio/backend/tests/test_openai_citation_markers_edge.py @@ -13,8 +13,8 @@ Reference: https://developers.openai.com/api/docs/guides/citation-formatting import importlib -# Streaming integration is exercised by ``_simulate_delta_stream`` further -# down, mirroring the head/buffer/flush dance from ``_stream_openai_responses``. +# Streaming is exercised by ``_simulate_delta_stream`` below, mirroring the +# head/buffer/flush dance from ``_stream_openai_responses``. _module = importlib.import_module("core.inference.external_provider") _replace_openai_citation_markers = _module._replace_openai_citation_markers _split_pending_citation_tail = _module._split_pending_citation_tail @@ -27,7 +27,7 @@ CITE_DELIM = "" def _marker(*source_ids: str, locator: str | None = None) -> str: """Build a ``\\ue200cite\\ue202[\\ue202...][\\ue202]\\ue201`` - marker. Accepts one or many ``source_ids`` plus an optional ``locator``.""" + marker from one or more ``source_ids`` plus an optional ``locator``.""" payload = f"{CITE_START}cite{CITE_DELIM}" + CITE_DELIM.join(source_ids) if locator: payload = f"{payload}{CITE_DELIM}{locator}" @@ -39,7 +39,7 @@ def _no_private_use(text: str) -> bool: # Harness mirroring the head/pending-tail/flush dance in -# `_stream_openai_responses`, so streaming tests skip the httpx mock. +# `_stream_openai_responses` so streaming tests skip the httpx mock. def _simulate_delta_stream( deltas: list[str], citations: list[dict], @@ -56,8 +56,8 @@ def _simulate_delta_stream( if head: emitted.append(head) if flush and pending: - # Mirror `_flush_pending_marker_tail`: drop the tail entirely if no - # closing stop byte arrived; the literal ``cite`` would leak otherwise. + # Mirror `_flush_pending_marker_tail`: drop the tail if no closing stop + # byte arrived; the literal ``cite`` would leak otherwise. if CITE_STOP not in pending: rendered = "" else: @@ -79,7 +79,7 @@ def _simulate_delta_stream( def test_multi_source_marker_all_resolve(): """\\ue200cite\\ue202id1\\ue202id2\\ue202id3\\ue201 expands to three links - when every id is known. Earlier regex captured only id1 and dropped id2/id3.""" + when every id is known. Earlier regex captured only id1.""" text = f"All three: {_marker('id1', 'id2', 'id3')}" citations = [ {"source_id": "id1", "url": "https://example.com/1"}, @@ -136,14 +136,14 @@ def test_marker_with_range_locator(): def test_marker_split_in_source_id(): - """Delta-1 ends mid-source-id (``\\ue200cite\\ue202tu``), delta-2 starts - with the rest (``rn0view0\\ue201``). The buffer stitches the halves - back together so they resolve to one link instead of leaking.""" + """Delta-1 ends mid-source-id (``\\ue200cite\\ue202tu``), delta-2 has the + rest (``rn0view0\\ue201``). The buffer stitches the halves so they resolve + to one link instead of leaking.""" full = f"See {_marker('turn0view0')} now." # Cut right after the second delim + "tu" inside the source id. cut = full.index("tu", full.index(CITE_START)) + len("tu") d1, d2 = full[:cut], full[cut:] - # Sanity check: delta-1 actually contains a partial marker. + # Sanity: delta-1 contains a partial marker. assert CITE_START in d1 and CITE_STOP not in d1 assert CITE_STOP in d2 citations = [{"source_id": "turn0view0", "url": "https://x"}] @@ -153,8 +153,8 @@ def test_marker_split_in_source_id(): def test_marker_split_at_start_byte(): - """Split exactly after the opening ``\\ue200`` byte; the buffer must - hold the lone open byte until the rest arrives.""" + """Split right after the opening ``\\ue200`` byte; the buffer must hold the + lone open byte until the rest arrives.""" full = f"Text {_marker('sid')} done" cut = full.index(CITE_START) + 1 # right AFTER the open byte d1, d2 = full[:cut], full[cut:] @@ -167,7 +167,7 @@ def test_marker_split_at_start_byte(): def test_marker_split_across_three_deltas(): """Worst case: marker chopped into three pieces across three deltas.""" full = f"A {_marker('threesplit')} B" - # cut at two points inside the marker + # Cut at two points inside the marker. open_pos = full.index(CITE_START) stop_pos = full.index(CITE_STOP) cut1 = open_pos + 4 @@ -206,21 +206,20 @@ def test_split_marker_unknown_source_is_dropped_cleanly(): def test_unterminated_marker_at_stream_end_dropped_on_flush(): - """Stream ends mid-marker (e.g. response.incomplete); the tail is - flushed with private-use bytes stripped, no `E202` text leaks.""" + """Stream ends mid-marker (e.g. response.incomplete); the flushed tail + strips private-use bytes, no `E202` text leaks.""" deltas = ["Some text ", f"{CITE_START}citetu", "rn0view0"] # no STOP ever out = _simulate_delta_stream(deltas, [], flush = True) assert _no_private_use(out) assert "E200" not in out and "E202" not in out - # Surrounding prose stays; we don't assert exact marker remainder. + # Surrounding prose stays; don't assert exact marker remainder. assert "Some text " in out def test_flush_resolves_marker_when_late_annotation_arrives(): - """Marker in a delta, matching annotation arrives later (on - response.output_text.annotation.added after the final delta). The - rewriter reads ``all_url_citations`` LIVE at flush, so the buffered - marker still resolves.""" + """Marker in a delta; the matching annotation arrives later (on + response.output_text.annotation.added after the final delta). The rewriter + reads ``all_url_citations`` LIVE at flush, so the buffered marker resolves.""" deltas = ["Look ", f"{CITE_START}cite{CITE_DELIM}late_sid"] pending = "" citations: list[dict] = [] @@ -291,8 +290,8 @@ def test_rewriter_idempotent_on_marker_free_text(): def test_only_marker_no_surrounding_text(): - """A delta that is JUST a marker (no prose) still renders correctly; - used to leak without the empty-string short-circuit in the split helper.""" + """A delta that is JUST a marker (no prose) renders correctly; it leaked + before the empty-string short-circuit in the split helper.""" text = _marker("solo") citations = [{"source_id": "solo", "url": "https://solo.example"}] out = _replace_openai_citation_markers(text, citations) @@ -311,18 +310,16 @@ def test_back_to_back_markers_with_no_separator(): def test_split_helper_buffers_only_after_last_open_byte(): - """A complete marker followed by an unterminated one: head includes - the complete marker, buffer holds only the trailing partial.""" + """Complete marker followed by an unterminated one: head includes the + complete marker, buffer holds only the trailing partial.""" complete = _marker("done") partial = f"{CITE_START}cite{CITE_DELIM}half" # no STOP text = f"pre {complete} mid {partial}" head, tail = _split_pending_citation_tail(text) assert head == f"pre {complete} mid " assert tail == partial - # And the head, once rewritten, drops every private-use byte. - rewritten = _replace_openai_citation_markers( - head, [{"source_id": "done", "url": "https://d"}] - ) + # Head, once rewritten, drops every private-use byte. + rewritten = _replace_openai_citation_markers(head, [{"source_id": "done", "url": "https://d"}]) assert rewritten == "pre [[1]](https://d) mid " @@ -357,7 +354,7 @@ def test_unknown_marker_does_not_perturb_citation_indexing(): {"source_id": "real_b", "url": "https://example.com/b"}, ] out = _replace_openai_citation_markers(text, citations) - # real_a is index 1; unknown does not take a slot. + # real_a is index 1; unknown takes no slot. assert "[[1]](https://example.com/a)" in out assert "[[2]](https://example.com/b)" in out assert _no_private_use(out) @@ -370,8 +367,8 @@ def test_unknown_marker_does_not_perturb_citation_indexing(): def test_unterminated_marker_does_not_leak_cite_residue(): - """Stream ends mid-marker: drop the whole tail rather than strip - codepoints and leave ``cite`` behind.""" + """Stream ends mid-marker: drop the whole tail rather than strip codepoints + and leave ``cite`` behind.""" half = f"Hi there {CITE_START}cite{CITE_DELIM}turn0view0" out = _simulate_delta_stream([half], [], flush = True) # Prose before the marker stays; no private-use bytes or cite residue. diff --git a/studio/backend/tests/test_openai_code_execution.py b/studio/backend/tests/test_openai_code_execution.py index 3d179371e3..63e94613ed 100644 --- a/studio/backend/tests/test_openai_code_execution.py +++ b/studio/backend/tests/test_openai_code_execution.py @@ -1,30 +1,13 @@ # 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 +"""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. +Covers: request body shaping (container_auto, container_reference), the cloud +guard (no shell tool on non-cloud base_urls), SSE translation of a +shell_call/shell_call_output pair into tool_start/tool_end events, container_id +surfacing as container_ready, and stale-container invalidation. """ import asyncio @@ -117,10 +100,7 @@ def test_shell_tool_added_on_cloud_with_container_auto(monkeypatch): _drive(run()) tools = captured["body"].get("tools") or [] - assert { - "type": "shell", - "environment": {"type": "container_auto"}, - } in tools + assert {"type": "shell", "environment": {"type": "container_auto"}} in tools def test_shell_tool_uses_container_reference_when_id_supplied(monkeypatch): @@ -195,8 +175,7 @@ def test_shell_tool_refused_for_non_cloud_base_url(monkeypatch): _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. + # Shell tool must not leak to local OpenAI-compat servers (they 400 on it). assert all(t.get("type") != "shell" for t in tools) @@ -269,7 +248,9 @@ def test_shell_call_emits_tool_start_and_end(monkeypatch): 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"} + # `_server_tool: True` marks a synthetic builtin so the frontend can tell + # hosted tools from user-declared functions on history replay. + assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la", "_server_tool": True} assert ends[0]["tool_call_id"] == "scall_1" assert "total 24" in ends[0]["result"] @@ -392,24 +373,22 @@ def test_stale_container_emits_invalidated(monkeypatch): def test_expired_container_triggers_transparent_retry(monkeypatch): - """When OpenAI 400s with 'Container is expired' on a request that - carried container_reference, the streamer retries once with the - container field stripped. The user never sees an error line — only - container_invalidated, then the normal stream from the retry. + """On a 'Container is expired' 400 for a container_reference request, the + streamer retries once with the container stripped; the user sees only + container_invalidated then the retry stream, never an error line. """ calls: list[dict] = [] def handler(request: httpx.Request) -> httpx.Response: body = json.loads(request.content.decode("utf-8")) calls.append(body) - # Find the shell tool entry to inspect environment.type. + # Inspect the shell tool's environment.type. shell_env_type = None for tool in body.get("tools", []) or []: if tool.get("type") == "shell": shell_env_type = tool.get("environment", {}).get("type") break - # First call carries container_reference -> 400 expired. - # Retry omits container -> normal SSE stream. + # container_reference -> 400 expired; retry omits container -> normal stream. if shell_env_type == "container_reference": return httpx.Response( 400, @@ -423,8 +402,8 @@ def test_expired_container_triggers_transparent_retry(monkeypatch): ).encode("utf-8"), headers = {"content-type": "application/json"}, ) - # Successful retry: minimal SSE — a completed response with a - # fresh container_id so container_ready latches. + # Successful retry: minimal SSE — completed response with a fresh + # container_id so container_ready latches. sse = _openai_sse( [ { @@ -460,8 +439,8 @@ def test_expired_container_triggers_transparent_retry(monkeypatch): lines = _drive(run()) events = _tool_events(lines) - # Two outbound HTTP calls were made: the expired-container attempt - # then the retry without the container field. + # Two outbound calls: the expired-container attempt, then the retry + # without the container field. assert len(calls) == 2 shell_types = [] for body in calls: @@ -470,14 +449,14 @@ def test_expired_container_triggers_transparent_retry(monkeypatch): shell_types.append(tool.get("environment", {}).get("type")) assert shell_types == ["container_reference", "container_auto"] - # container_invalidated emitted (frontend will null its stored id). + # container_invalidated emitted (frontend nulls its stored id). assert any(e.get("type") == "container_invalidated" for e in events) # container_ready emitted from the retry stream with the fresh id. assert any( e.get("type") == "container_ready" and e.get("container_id") == "cntr_fresh_111" for e in events ) - # CRUCIALLY: no SSE error line surfaced to the chat — only completion. + # CRUCIALLY: no SSE error line surfaced to the chat. error_lines = [ line for line in lines @@ -487,8 +466,8 @@ def test_expired_container_triggers_transparent_retry(monkeypatch): def test_expired_container_retries_only_once(monkeypatch): - """If the retry ALSO fails (any 4xx, expired or otherwise), the - error is surfaced normally — no infinite retry loop. + """If the retry ALSO fails (any 4xx), the error surfaces normally — + no infinite retry loop. """ call_count = {"n": 0} @@ -527,11 +506,8 @@ def test_expired_container_retries_only_once(monkeypatch): lines = _drive(run()) - # Exactly two calls (first + one retry). Third would mean an - # infinite loop. + # Exactly two calls (first + one retry); a third would be a loop. assert call_count["n"] == 2 # The second failure surfaces normally as an error SSE line. - error_lines = [ - line for line in lines if '"error"' in line and "_toolEvent" not in line - ] + error_lines = [line for line in lines if '"error"' in line and "_toolEvent" not in line] assert len(error_lines) >= 1 diff --git a/studio/backend/tests/test_openai_compaction.py b/studio/backend/tests/test_openai_compaction.py index f0599952c6..c7de0a9aed 100644 --- a/studio/backend/tests/test_openai_compaction.py +++ b/studio/backend/tests/test_openai_compaction.py @@ -4,14 +4,13 @@ """Unit tests for OpenAI Responses API context_management wiring. OpenAI's Responses API supports server-side compaction via -``context_management: [{type:"compaction", compact_threshold:N}]``. -There is no beta header and no dated version pin; the threshold is -silently accepted and the API runs the compaction step when the -rendered prompt crosses it. +``context_management: [{type:"compaction", compact_threshold:N}]``. No +beta header, no dated version pin; the threshold is silently accepted and +compaction runs when the rendered prompt crosses it. -These tests pin: the body shape when threshold is set on cloud OpenAI, -the silent no-op when the base URL is non-cloud, and the -omitted-threshold pass-through. +These pin: the body shape when threshold is set on cloud OpenAI, the +silent no-op on non-cloud base URLs, and the omitted-threshold +pass-through. """ import asyncio @@ -32,8 +31,7 @@ def _capture(monkeypatch, *, base_url: str, threshold) -> dict: def handler(request: httpx.Request) -> httpx.Response: captured["body"] = json.loads(request.content.decode("utf-8")) - # Send an empty Responses-shaped SSE stream so the helper exits - # cleanly. + # Empty Responses-shaped SSE stream so the helper exits cleanly. return httpx.Response( 200, content = ( @@ -88,8 +86,8 @@ def test_cloud_openai_sets_compaction_block(monkeypatch): def test_cloud_openai_below_default_threshold_passes_through(monkeypatch): - # Studio doesn't clamp the OpenAI side -- the API accepts whatever - # the caller sends, so a small probe like 60k still goes through. + # Studio doesn't clamp the OpenAI side -- the API accepts whatever the + # caller sends, so a small probe like 60k still goes through. captured = _capture( monkeypatch, base_url = "https://api.openai.com/v1", @@ -105,8 +103,8 @@ def test_cloud_openai_below_default_threshold_passes_through(monkeypatch): def test_non_cloud_base_silently_drops_compaction(monkeypatch): # ollama / llama.cpp / "custom" presets collapse to provider="openai" - # but don't implement context_management. Sending the field would - # 400 those servers, so it must NOT appear on the wire. + # but lack context_management. Sending the field would 400 them, so it + # must NOT appear on the wire. captured = _capture( monkeypatch, base_url = "http://127.0.0.1:11434/v1", @@ -120,9 +118,9 @@ def test_non_cloud_base_silently_drops_compaction(monkeypatch): def test_azure_openai_base_url_carries_compaction_block(monkeypatch): # Azure OpenAI Foundry exposes the same /v1/responses extensions - # (context_management, prompt_cache_retention, container shell) - # under a *.openai.azure.com base URL. Treat it as cloud so the - # compaction field actually reaches the API. + # (context_management, prompt_cache_retention, container shell) under + # a *.openai.azure.com base URL. Treat it as cloud so the compaction + # field reaches the API. captured = _capture( monkeypatch, base_url = "https://my-resource.openai.azure.com/openai/v1", @@ -132,14 +130,14 @@ def test_azure_openai_base_url_carries_compaction_block(monkeypatch): {"type": "compaction", "compact_threshold": 200_000} ] # Sibling Azure-cloud extension: prompt_cache_retention should also - # be set so caching works the same way on Azure deployments. + # be set so caching works the same on Azure deployments. assert captured["body"].get("prompt_cache_retention") == "24h" def test_azure_openai_mixed_case_base_url_matches(monkeypatch): - # The match is case-insensitive so URLs copy-pasted from the Azure - # portal (which sometimes capitalise the resource name) still get - # the cloud-only fields. + # Case-insensitive match so URLs copy-pasted from the Azure portal + # (which sometimes capitalise the resource name) still get the + # cloud-only fields. captured = _capture( monkeypatch, base_url = "https://My-Resource.OpenAI.Azure.Com/openai/v1", @@ -151,12 +149,11 @@ def test_azure_openai_mixed_case_base_url_matches(monkeypatch): def test_cloud_gate_uses_hostname_not_substring(monkeypatch): - # CodeQL py/incomplete-url-substring-sanitization: an attacker who - # controls the configured base_url could embed `api.openai.com` or - # `.openai.azure.com` as part of a path or a subdomain on an - # arbitrary host to slip the cloud-only request body fields to a - # server they control. The hostname-anchored helper must reject - # both shapes. + # CodeQL py/incomplete-url-substring-sanitization: an attacker + # controlling base_url could embed `api.openai.com` or + # `.openai.azure.com` in a path or subdomain on an arbitrary host to + # slip cloud-only body fields to their own server. The + # hostname-anchored helper must reject both shapes. for evil in [ "https://evil.com/api.openai.com/v1", "https://api.openai.com.attacker.com/v1", @@ -188,19 +185,17 @@ def test_omitted_threshold_no_body_field(monkeypatch): def test_chat_completion_request_accepts_any_positive_compaction_threshold(): - # Codex follow-up: the field is documented as a no-op for non-cloud - # OpenAI bases and every non-OpenAI provider, so a cross-provider - # schema floor would 422 perfectly valid Anthropic / ollama / - # llama.cpp requests that happen to carry the field. Keep schema - # floor at ge=1 (any positive int) and rely on per-provider - # helpers (_stream_openai_responses / _stream_anthropic) to - # enforce or clamp the real floor. + # Codex follow-up: the field is a no-op for non-cloud OpenAI bases and + # every non-OpenAI provider, so a cross-provider schema floor would + # 422 valid Anthropic / ollama / llama.cpp requests carrying it. Keep + # the schema floor at ge=1 (any positive int) and let per-provider + # helpers (_stream_openai_responses / _stream_anthropic) enforce or + # clamp the real floor. import pytest as _pytest from models.inference import ChatCompletionRequest - # Non-positive values still rejected so blank-string posts don't - # sneak through. + # Non-positive values rejected so blank-string posts don't sneak in. with _pytest.raises(Exception): ChatCompletionRequest.model_validate( { @@ -211,10 +206,10 @@ def test_chat_completion_request_accepts_any_positive_compaction_threshold(): ) # Any positive int passes schema validation, including values that - # would be no-ops on the OpenAI cloud path. This is intentional -- - # the OpenAI helper drops the field on non-cloud bases and - # forwards-as-is on cloud bases; if the value is below the model's - # effective floor, the upstream API surfaces the error. + # are no-ops on the OpenAI cloud path. Intentional -- the OpenAI + # helper drops the field on non-cloud bases and forwards as-is on + # cloud bases; if it's below the model's effective floor, the upstream + # API surfaces the error. for v in (1, 5_000, 9_999, 10_000, 200_000): req = ChatCompletionRequest.model_validate( { diff --git a/studio/backend/tests/test_openai_container_crud.py b/studio/backend/tests/test_openai_container_crud.py index 161a6fab83..e0604527fd 100644 --- a/studio/backend/tests/test_openai_container_crud.py +++ b/studio/backend/tests/test_openai_container_crud.py @@ -4,11 +4,11 @@ """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. +- list / create / delete all send ``OpenAI-Beta: containers=v1``. Without + it, OpenAI silently no-ops the DELETE but still returns 200 + ``{"deleted": true}``. +- ``delete_openai_container`` raises when the body omits + ``{"deleted": true}``, even on a 2xx response. """ from __future__ import annotations @@ -28,11 +28,10 @@ def _drive(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.""" + """Wire `handler` for the shared `_http_client` AND any per-call + `httpx.AsyncClient(...)`. delete_openai_container creates a fresh + AsyncClient (see external_provider.delete_openai_container), so we + must also intercept that constructor.""" transport = httpx.MockTransport(handler) monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) real_async_client = httpx.AsyncClient @@ -80,17 +79,12 @@ def test_create_sends_openai_beta_header(monkeypatch): 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) - ) + 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, - } + assert seen["body"]["expires_after"] == {"anchor": "last_active_at", "minutes": 30} def test_delete_sends_openai_beta_header_and_accepts_confirmation(monkeypatch): @@ -115,13 +109,12 @@ def test_delete_sends_openai_beta_header_and_accepts_confirmation(monkeypatch): 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.""" + silently rejected (e.g. before we sent OpenAI-Beta). Guard: when the + body omits ``deleted: true``, surface an error so the UI reports the + failure instead of false success.""" def handler(request: httpx.Request) -> httpx.Response: - # 200 but no deleted flag — simulate an unexpected payload shape. + # 200 but no deleted flag — unexpected payload shape. return httpx.Response(200, json = {"id": "cntr_x", "object": "container"}) _mock_http_client(monkeypatch, handler) @@ -164,10 +157,9 @@ def test_delete_propagates_openai_4xx(monkeypatch): 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.""" + """OpenAI keeps containers in /v1/containers with status="expired" + after their idle TTL passes — unusable but still listed. The list + route must drop them so the picker shows only usable containers.""" from routes import inference as inf_mod from models.inference import OpenAIContainerRequest diff --git a/studio/backend/tests/test_openai_image_generation.py b/studio/backend/tests/test_openai_image_generation.py index 6d415316d4..ace57588d3 100644 --- a/studio/backend/tests/test_openai_image_generation.py +++ b/studio/backend/tests/test_openai_image_generation.py @@ -3,18 +3,11 @@ """Unit tests for OpenAI Responses API image_generation tool wiring. -The image_generation tool is a server-side Responses-API tool: -``{type: "image_generation"}`` in the request's tools array, and the -result comes back as an ``image_generation_call`` output item carrying -the base64 image on ``result``. Studio translates the output item -into ``_toolEvent`` chunks (``tool_start`` with `kind:"image"`, -``tool_end`` with ``image_b64`` + ``image_mime``) so the chat adapter -can render the image inline. - -These tests pin: the tool is added to the outbound body only when the -caller asks for it on a cloud OpenAI base; the SSE output_item.done -for ``image_generation_call`` produces the expected _toolEvent chunks; -non-cloud bases drop the tool silently. +The tool is a server-side Responses-API tool (``{type: "image_generation"}``); +the result comes back as an ``image_generation_call`` output item, which Studio +translates into ``_toolEvent`` chunks so the chat adapter renders it inline. +Tests pin: the tool is added to the body only on a cloud OpenAI base when asked +for, the done event produces the expected chunks, and non-cloud bases drop it. """ import asyncio @@ -75,8 +68,8 @@ def _capture_body(monkeypatch, *, base_url: str, enabled_tools) -> dict: def _collect_tool_events(monkeypatch) -> list[dict]: - """Drive a Responses stream that emits one image_generation_call done - event and return the parsed _toolEvent chunks.""" + """Drive a Responses stream with one image_generation_call done event and + return the parsed _toolEvent chunks.""" sse = ( b"event: response.output_item.done\n" @@ -207,9 +200,12 @@ def test_image_generation_done_emits_tool_event_chunks(monkeypatch): ends = [e for e in image_events if e.get("type") == "tool_end"] assert len(starts) == 1, image_events assert len(ends) == 1, image_events + # `_server_tool: True` marks this as a provider-side synthetic tool card + # for the frontend's history serializer. assert starts[0]["arguments"] == { "kind": "image", "prompt": "A photorealistic cat sitting", + "_server_tool": True, "openai_image_generation_call_id": "img_abc", } assert ends[0]["image_b64"] == "AAAA" diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py index 22ccba7058..f7d7e83a43 100644 --- a/studio/backend/tests/test_openai_responses_translation.py +++ b/studio/backend/tests/test_openai_responses_translation.py @@ -5,14 +5,14 @@ Unit tests for the OpenAI `/v1/responses` translation in external_provider. Covers: -- Request body shape: system messages collapse into `instructions`, user/ - assistant messages go into `input`, sampling knobs Responses does not - support (presence_penalty, top_k) are not forwarded. -- SSE translation: `response.output_text.delta` events become OpenAI Chat - Completions chunks, `response.completed` emits a `finish_reason: stop` - chunk, the stream terminates with `data: [DONE]`. -- Image parts in user content are rewritten from Chat Completions - `{type: image_url, image_url: {url}}` into Responses +- Request body shape: system messages collapse into `instructions`, + user/assistant messages go into `input`, and unsupported sampling knobs + (presence_penalty, top_k) are not forwarded. +- SSE translation: `response.output_text.delta` → Chat Completions chunks, + `response.completed` → a `finish_reason: stop` chunk, stream ends with + `data: [DONE]`. +- Image parts rewritten from Chat Completions + `{type: image_url, image_url: {url}}` to Responses `{type: input_image, image_url: }`. """ @@ -101,9 +101,9 @@ def test_responses_request_body_uses_input_and_instructions(monkeypatch): assert body["input"] == [{"role": "user", "content": "Hi"}] assert body["max_output_tokens"] == 512 assert body["stream"] is True - # Responses API on reasoning-class models (gpt-5.x / o3 / gpt-4.5 — the - # only OpenAI ids the registry allowlist exposes) rejects these as - # `Unsupported parameter`. Make sure we never silently forward them. + # Responses API on reasoning-class models (gpt-5.x / o3 / gpt-4.5 — the only + # OpenAI ids the registry allowlist exposes) rejects these as `Unsupported + # parameter`. Never silently forward them. assert "temperature" not in body assert "top_p" not in body assert "presence_penalty" not in body @@ -154,10 +154,7 @@ def test_responses_translates_image_parts(monkeypatch): parts = captured["body"]["input"][0]["content"] assert parts[0] == {"type": "input_text", "text": "What is this?"} - assert parts[1] == { - "type": "input_image", - "image_url": "data:image/png;base64,AAA", - } + assert parts[1] == {"type": "input_image", "image_url": "data:image/png;base64,AAA"} # No max_output_tokens key when caller passes max_tokens=None. assert "max_output_tokens" not in captured["body"] @@ -196,7 +193,7 @@ def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch): lines = _drive(run()) - # Drop empty / non-data lines for assertion clarity. + # Keep only data lines for assertion clarity. data_lines = [line for line in lines if line.startswith("data:")] payloads = [] for line in data_lines: @@ -215,6 +212,254 @@ def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch): assert payloads[-1] == "[DONE]" +def test_responses_function_call_output_translates_to_delta_tool_calls(monkeypatch): + """Round 12: function tools forwarded into /v1/responses must have their + `function_call` output items translated back into Chat Completions + delta.tool_calls, and the terminal chunk must emit + finish_reason="tool_calls" (not "stop") so the frontend's accumulator runs + the function.""" + + def handler(request: httpx.Request) -> httpx.Response: + events = [ + {"type": "response.created"}, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_abc", + "call_id": "call_xyz", + "name": "get_weather", + "arguments": '{"city":"SF"}', + }, + }, + {"type": "response.completed", "response": {}}, + ] + return httpx.Response( + 200, + content = _responses_sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + lines = await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "weather?"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } + ], + ) + ) + await client.close() + return lines + + lines = _drive(run()) + payloads = [ + json.loads(line[len("data:") :].strip()) + for line in lines + if line.startswith("data:") and line[len("data:") :].strip() != "[DONE]" + ] + tool_call_deltas = [ + p + for p in payloads + if isinstance(p, dict) + and p.get("choices") + and p["choices"][0].get("delta", {}).get("tool_calls") + ] + assert tool_call_deltas, payloads + tc = tool_call_deltas[0]["choices"][0]["delta"]["tool_calls"][0] + assert tc["id"] == "call_xyz" + assert tc["function"]["name"] == "get_weather" + assert tc["function"]["arguments"] == '{"city":"SF"}' + # Final chunk reports tool_calls, not stop. + terminal = next( + p + for p in payloads + if isinstance(p, dict) + and p.get("choices") + and p["choices"][0].get("finish_reason") in ("stop", "tool_calls") + ) + assert terminal["choices"][0]["finish_reason"] == "tool_calls", payloads + + +def test_responses_parallel_function_calls_get_distinct_indices(monkeypatch): + """Round 13: parallel function_call items must land on distinct + delta.tool_calls[].index slots so index-keyed clients don't collapse the + second call into the first.""" + + def handler(request: httpx.Request) -> httpx.Response: + events = [ + {"type": "response.created"}, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_a", + "call_id": "call_a", + "name": "lookup_a", + "arguments": "{}", + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_b", + "call_id": "call_b", + "name": "lookup_b", + "arguments": "{}", + }, + }, + {"type": "response.completed", "response": {}}, + ] + return httpx.Response( + 200, + content = _responses_sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + lines = await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "x"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + tools = [ + { + "type": "function", + "function": { + "name": "lookup_a", + "parameters": {"type": "object"}, + }, + }, + { + "type": "function", + "function": { + "name": "lookup_b", + "parameters": {"type": "object"}, + }, + }, + ], + ) + ) + await client.close() + return lines + + lines = _drive(run()) + indices: list[int] = [] + for raw in lines: + if not raw.startswith("data:"): + continue + payload = raw[len("data:") :].strip() + if payload == "[DONE]": + continue + try: + obj = json.loads(payload) + except Exception: + continue + delta = (obj.get("choices") or [{}])[0].get("delta") or {} + for tc in delta.get("tool_calls") or []: + indices.append(tc.get("index")) + assert indices == [0, 1], indices + + +def test_responses_follow_up_tool_result_uses_function_call_output_items(monkeypatch): + """Round 13: a second turn after a Responses function call must serialize + the tool_calls history and tool result as Responses `function_call` / + `function_call_output` input items, not Chat Completions role="tool" + content.""" + 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.created"}, + {"type": "response.completed", "response": {}}, + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + await _collect( + client._stream_openai_responses( + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_xyz", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"SF"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_xyz", + "content": "sunny", + }, + {"role": "user", "content": "thanks"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + ) + ) + await client.close() + + _drive(run()) + items = captured["body"]["input"] + types = [it.get("type") or it.get("role") for it in items] + assert "function_call" in types, items + assert "function_call_output" in types, items + fc = next(it for it in items if it.get("type") == "function_call") + assert fc["call_id"] == "call_xyz" + assert fc["name"] == "get_weather" + assert fc["arguments"] == '{"city":"SF"}' + fco = next(it for it in items if it.get("type") == "function_call_output") + assert fco["call_id"] == "call_xyz" + assert fco["output"] == "sunny" + + def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch): def handler(request: httpx.Request) -> httpx.Response: events = [ @@ -249,8 +494,7 @@ def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch) finish_reasons = [ json.loads(line[len("data:") :].strip())["choices"][0]["finish_reason"] for line in lines - if line.startswith("data:") - and line[len("data:") :].strip() not in ("", "[DONE]") + if line.startswith("data:") and line[len("data:") :].strip() not in ("", "[DONE]") ] assert "length" in finish_reasons @@ -482,8 +726,7 @@ def test_responses_reasoning_summary_wrapped_in_think_tags(monkeypatch): data_lines = [ line[len("data:") :].strip() for line in lines - if line.startswith("data:") - and line[len("data:") :].strip() not in ("", "[DONE]") + if line.startswith("data:") and line[len("data:") :].strip() not in ("", "[DONE]") ] payloads = [json.loads(raw) for raw in data_lines] combined = "".join( diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 84f3e41998..1d994b46c0 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1,25 +1,19 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -""" -Tests for the OpenAI /v1/chat/completions client-side tool pass-through. +"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through. -Covers: -- ChatCompletionRequest accepts standard OpenAI `tools` / `tool_choice` / `stop`. -- ChatMessage accepts role="tool" with `tool_call_id` and role="assistant" - with `content: None` + `tool_calls`. -- ChatCompletionRequest carries unknown fields via `extra="allow"`. -- anthropic_tool_choice_to_openai() covers all four Anthropic shapes. -- _build_passthrough_payload() honors a caller-supplied tool_choice and - defaults to "auto" when unset. -- _friendly_error() maps httpx transport errors to a "Lost connection" - message so passthrough failures are legible instead of bare 500s. - -No running server or GPU required. +Covers ChatMessage tool/assistant roles, ChatCompletionRequest tool fields and +extra="allow", anthropic_tool_choice_to_openai, _build_passthrough_payload +tool_choice propagation, and _friendly_error's httpx-to-"Lost connection" +mapping. No server or GPU required. """ import os import sys +import asyncio +import json +from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") sys.path.insert(0, _backend) @@ -32,11 +26,24 @@ from pydantic import ValidationError from models.inference import ( ChatCompletionRequest, ChatMessage, + CompletionChoice, + CompletionMessage, ) from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) -from routes.inference import _build_passthrough_payload, _friendly_error +from routes.inference import ( + _build_openai_passthrough_body, + _build_passthrough_payload, + _clamp_finish_reason, + _effective_max_tokens, + _extract_content_parts, + _friendly_error, + _openai_stream_usage_chunk, + _set_or_prepend_system_message, + openai_chat_completions, +) +from state.tool_policy import reset_tool_policy # ===================================================================== @@ -112,8 +119,8 @@ class TestChatMessageToolRoles: ChatMessage(role = "function", content = "x") def test_content_absent_on_assistant_tool_call_defaults_to_none(self): - # Assistant messages that carry only tool_calls are the one - # documented case where `content=None` is permitted. + # Assistant messages carrying only tool_calls are the one documented + # case where `content=None` is permitted. msg = ChatMessage( role = "assistant", tool_calls = [ @@ -128,9 +135,9 @@ class TestChatMessageToolRoles: def test_tool_role_missing_tool_call_id_left_for_request_validator(self): # Per-message: missing tool_call_id is now allowed at this layer. - # ChatCompletionRequest's walkback fills it in from the prior - # assistant tool_calls; see test_inference_model_validation.py for - # the resolution coverage. + # ChatCompletionRequest's walkback fills it from the prior assistant + # tool_calls; see test_inference_model_validation.py for resolution + # coverage. msg = ChatMessage(role = "tool", content = '{"temperature": 72}') assert msg.tool_call_id is None assert msg.content == '{"temperature": 72}' @@ -159,13 +166,14 @@ class TestChatMessageToolRoles: with pytest.raises(ValidationError): ChatMessage(role = "user", content = []) - def test_tool_empty_content_rejected(self): - with pytest.raises(ValidationError) as exc_info: - ChatMessage(role = "tool", tool_call_id = "call_1", content = "") - assert "content" in str(exc_info.value) + def test_tool_empty_content_accepted(self): + # Empty tool output (mkdir, git add, ...) is routine in agentic loops; + # OpenAI and llama-server both accept it, so Studio must not 400. + msg = ChatMessage(role = "tool", tool_call_id = "call_1", content = "") + assert msg.content == "" def test_assistant_without_content_or_tool_calls_tolerated(self): - # Stop-button leaves an empty assistant turn; tolerate so replay round-trips. + # Stop-button leaves an empty assistant turn; tolerate for replay. msg = ChatMessage(role = "assistant") assert msg.content is None assert msg.tool_calls is None @@ -263,10 +271,7 @@ class TestChatCompletionRequestToolFields: assert self._make(stop = "\nUser:").stop == "\nUser:" def test_stop_list(self): - assert self._make(stop = ["\nUser:", "\nAssistant:"]).stop == [ - "\nUser:", - "\nAssistant:", - ] + assert self._make(stop = ["\nUser:", "\nAssistant:"]).stop == ["\nUser:", "\nAssistant:"] def test_tools_default_none(self): req = self._make() @@ -275,18 +280,19 @@ class TestChatCompletionRequestToolFields: assert req.stop is None def test_extra_fields_accepted(self): - # `frequency_penalty`, `seed`, `response_format` are not yet - # explicitly declared but must survive Pydantic parsing now that - # extra="allow" is set. + # `frequency_penalty` and `response_format` are not yet explicitly + # declared but must survive Pydantic parsing now that extra="allow" is + # set. `seed` is declared and should land on the typed field instead. req = self._make( frequency_penalty = 0.5, seed = 42, response_format = {"type": "json_object"}, ) + assert req.seed == 42 # Extras land in model_extra assert req.model_extra is not None assert req.model_extra.get("frequency_penalty") == 0.5 - assert req.model_extra.get("seed") == 42 + assert "seed" not in req.model_extra assert req.model_extra.get("response_format") == {"type": "json_object"} def test_unsloth_extensions_still_work(self): @@ -300,27 +306,16 @@ class TestChatCompletionRequestToolFields: assert req.session_id == "abc" def test_stream_defaults_false_matching_openai_spec(self): - # OpenAI's /v1/chat/completions spec defaults `stream` to false. - # Studio previously defaulted to true, which broke naive curl - # clients (and .NET / System.Text.Json SDKs per #5047) that omit - # `stream` -- they expect a JSON blob, got SSE. - # Pin the corrected default so it can't silently regress. + # OpenAI defaults `stream` to false. Studio used to default true, + # breaking naive curl/.NET clients (#5047) that omit it. Pin the fix. req = self._make() assert req.stream is False - def test_post_without_stream_field_decodes_to_stream_false_over_http( - self, monkeypatch - ): - # Wire-level guard for the same default: a POST body that omits - # `stream` entirely (the exact shape naive curl / .NET clients - # send) must deserialise into stream=False *and* the response - # must be `application/json`, never `text/event-stream`. - # Mounts the real `routes.inference.router` so this catches - # regressions in middleware/aliasing on the actual endpoint - # (e.g. someone adding a request layer that injects stream=True - # before pydantic builds the model). Backends are bypassed by - # routing through `provider_type` and stubbing the external - # provider proxy. + def test_post_without_stream_field_decodes_to_stream_false_over_http(self, monkeypatch): + # Wire-level guard: a POST body omitting `stream` must deserialise to + # stream=False and return application/json, never text/event-stream. + # Mounts the real router to catch middleware/aliasing regressions; + # backends are bypassed via provider_type + a stubbed proxy. from fastapi import FastAPI from fastapi.responses import JSONResponse from fastapi.testclient import TestClient @@ -353,6 +348,153 @@ class TestChatCompletionRequestToolFields: assert "text/event-stream" not in resp.headers["content-type"] assert captured["stream"] is False + def _v1_client( + self, + monkeypatch, + llama_backend, + inference_backend = None, + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + import routes.inference as inference_route + from auth.authentication import get_current_subject + from utils.api_errors import install_api_error_handlers + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama_backend) + if inference_backend is not None: + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: inference_backend) + + app = FastAPI() + app.include_router(inference_route.router, prefix = "/v1") + install_api_error_handlers(app) + app.dependency_overrides[get_current_subject] = lambda: "test-user" + return TestClient(app) + + def _assert_unsupported_param(self, response, param): + assert response.status_code == 400 + body = response.json() + assert body["error"]["param"] == param + assert body["error"]["code"] == "unsupported_parameter" + + def _assert_unsupported_n(self, response): + self._assert_unsupported_param(response, "n") + + def test_n_allows_openai_chat_completion_range(self): + req = self._make(n = 128) + assert req.n == 128 + with pytest.raises(ValidationError): + self._make(n = 129) + + def test_n_rejected_for_external_provider_path(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "n": 2, + }, + ) + self._assert_unsupported_n(resp) + + def test_logprobs_rejected_until_supported(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "logprobs": True, + }, + ) + self._assert_unsupported_param(resp, "logprobs") + + def test_top_logprobs_rejected_until_supported(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "top_logprobs": 3, + }, + ) + self._assert_unsupported_param(resp, "top_logprobs") + + def test_n_rejected_for_gguf_streaming_path(self, monkeypatch): + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "n": 2, + }, + ) + self._assert_unsupported_n(resp) + + def test_n_rejected_for_gguf_tools_passthrough_path(self, monkeypatch): + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = True + is_vision = False + _is_audio = False + + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + "n": 2, + }, + ) + self._assert_unsupported_n(resp) + + def test_n_rejected_for_non_gguf_path(self, monkeypatch): + class _NoGGUFBackend: + is_loaded = False + + class _InferenceBackend: + active_model_name = "test-model" + models = {"test-model": {}} + + client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "n": 2, + }, + ) + self._assert_unsupported_n(resp) + def test_multiturn_tool_loop_messages(self): req = ChatCompletionRequest( messages = [ @@ -411,13 +553,8 @@ class TestAnthropicToolChoiceToOpenAI: assert anthropic_tool_choice_to_openai({"type": "none"}) == "none" def test_tool_named(self): - result = anthropic_tool_choice_to_openai( - {"type": "tool", "name": "get_weather"} - ) - assert result == { - "type": "function", - "function": {"name": "get_weather"}, - } + result = anthropic_tool_choice_to_openai({"type": "tool", "name": "get_weather"}) + assert result == {"type": "function", "function": {"name": "get_weather"}} def test_tool_missing_name_returns_none(self): assert anthropic_tool_choice_to_openai({"type": "tool"}) is None @@ -470,17 +607,226 @@ class TestBuildPassthroughPayloadToolChoice: body = _build_passthrough_payload(**self._args(), tool_choice = tc) assert body["tool_choice"] == tc - def test_stream_adds_include_usage(self): + def test_stream_omits_usage_options_when_client_did_not_request_them(self): args = self._args() args["stream"] = True body = _build_passthrough_payload(**args) + assert "stream_options" not in body + + def test_stream_forwards_include_usage_when_client_requests_it(self): + args = self._args() + args["stream"] = True + body = _build_passthrough_payload( + **args, + stream_options = {"include_usage": True}, + ) assert body.get("stream_options") == {"include_usage": True} + def test_stream_forwards_include_usage_false_when_client_requests_it(self): + args = self._args() + args["stream"] = True + body = _build_passthrough_payload( + **args, + stream_options = {"include_usage": False}, + ) + assert body.get("stream_options") == {"include_usage": False} + def test_repetition_penalty_renamed(self): body = _build_passthrough_payload(**self._args(), repetition_penalty = 1.1) assert body.get("repeat_penalty") == 1.1 assert "repetition_penalty" not in body + def test_passthrough_body_merges_system_and_developer_messages(self): + payload = ChatCompletionRequest( + model = "default", + messages = [ + {"role": "system", "content": "original system"}, + {"role": "developer", "content": "developer rules"}, + {"role": "user", "content": "hi"}, + ], + tools = self._args()["openai_tools"], + ) + + body = _build_openai_passthrough_body(payload, backend_ctx = 4096) + + assert body["messages"] == [ + {"role": "system", "content": "original system\n\ndeveloper rules"}, + {"role": "user", "content": "hi"}, + ] + + +# ===================================================================== +# Passthrough reasoning kwargs — enable_thinking / reasoning_effort / +# preserve_thinking must reach llama-server via chat_template_kwargs, +# gated on template capabilities like the non-passthrough paths. +# ===================================================================== + + +def _reasoning_backend( + supports_reasoning = True, + reasoning_style = "enable_thinking", + reasoning_always_on = False, + supports_preserve_thinking = False, +): + """Bare LlamaCppBackend with just the reasoning capability flags set, + so _build_openai_passthrough_body exercises the real + _request_reasoning_kwargs gating.""" + from core.inference.llama_cpp import LlamaCppBackend + + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._supports_reasoning = supports_reasoning + backend._reasoning_style = reasoning_style + backend._reasoning_always_on = reasoning_always_on + backend._supports_preserve_thinking = supports_preserve_thinking + return backend + + +class TestPassthroughReasoningKwargs: + def _payload(self, **fields): + return ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + **fields, + ) + + def test_enable_thinking_forwarded(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = False), + backend_ctx = 4096, + llama_backend = _reasoning_backend(), + ) + assert body["chat_template_kwargs"] == {"enable_thinking": False} + + def test_preserve_thinking_forwarded_when_template_supports_it(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = True, preserve_thinking = True), + backend_ctx = 4096, + llama_backend = _reasoning_backend(supports_preserve_thinking = True), + ) + assert body["chat_template_kwargs"] == { + "enable_thinking": True, + "preserve_thinking": True, + } + + def test_preserve_thinking_dropped_when_template_lacks_it(self): + body = _build_openai_passthrough_body( + self._payload(preserve_thinking = True), + backend_ctx = 4096, + llama_backend = _reasoning_backend(supports_preserve_thinking = False), + ) + assert "chat_template_kwargs" not in body + + def test_reasoning_effort_forwarded_for_effort_style_models(self): + body = _build_openai_passthrough_body( + self._payload(reasoning_effort = "high"), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"), + ) + assert body["chat_template_kwargs"] == {"reasoning_effort": "high"} + + def test_enable_thinking_maps_to_effort_for_effort_style_models(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = False), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"), + ) + assert body["chat_template_kwargs"] == {"reasoning_effort": "low"} + + def test_always_on_reasoning_skips_thinking_kwargs(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = False), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_always_on = True), + ) + assert "chat_template_kwargs" not in body + + def test_no_reasoning_fields_omits_chat_template_kwargs(self): + body = _build_openai_passthrough_body( + self._payload(), + backend_ctx = 4096, + llama_backend = _reasoning_backend(supports_preserve_thinking = True), + ) + assert "chat_template_kwargs" not in body + + +# ===================================================================== +# OpenAI API compatibility helpers — verified spec edge cases +# ===================================================================== + + +class TestOpenAICompatibilityHelpers: + def test_max_completion_tokens_wins_over_deprecated_max_tokens(self): + payload = SimpleNamespace(max_tokens = 128, max_completion_tokens = 64) + assert _effective_max_tokens(payload) == 64 + + @pytest.mark.parametrize( + "finish_reason", + ["stop", "length", "tool_calls", "content_filter", "function_call"], + ) + def test_clamp_finish_reason_preserves_openai_finish_reasons(self, finish_reason): + assert _clamp_finish_reason(finish_reason) == finish_reason + + def test_clamp_finish_reason_defaults_unknown_to_stop(self): + assert _clamp_finish_reason(None) == "stop" + assert _clamp_finish_reason("unexpected") == "stop" + + def test_non_streaming_completion_choice_accepts_tool_calls_finish_reason(self): + choice = CompletionChoice( + index = 0, + message = CompletionMessage(content = ""), + finish_reason = "tool_calls", + ) + assert choice.finish_reason == "tool_calls" + + def test_stream_usage_chunk_requires_include_usage(self): + usage = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} + payload = SimpleNamespace(stream_options = None) + assert ( + _openai_stream_usage_chunk(payload, "chatcmpl-test", 123, "model", usage, None) is None + ) + + payload.stream_options = {"include_usage": True} + line = _openai_stream_usage_chunk(payload, "chatcmpl-test", 123, "model", usage, None) + assert line is not None + assert '"choices":[]' in line + assert '"usage"' in line + + def test_stream_usage_chunk_coerces_nullable_counts(self): + payload = SimpleNamespace(stream_options = {"include_usage": True}) + line = _openai_stream_usage_chunk( + payload, + "chatcmpl-test", + 123, + "model", + {"prompt_tokens": None, "completion_tokens": 7, "total_tokens": None}, + None, + ) + + assert line is not None + parsed = json.loads(line.removeprefix("data: ")) + usage = parsed["usage"] + assert usage["prompt_tokens"] == 0 + assert usage["completion_tokens"] == 7 + assert usage["total_tokens"] == 7 + + def test_developer_message_preserves_existing_system_prompt(self): + payload = ChatCompletionRequest( + messages = [ + {"role": "system", "content": "original system"}, + {"role": "developer", "content": "developer rules"}, + {"role": "user", "content": "hi"}, + ] + ) + for message in payload.messages: + if message.role == "developer": + message.role = "system" + + system_prompt, chat_messages, image_b64 = _extract_content_parts(payload.messages) + + assert system_prompt == "original system\n\ndeveloper rules" + assert chat_messages == [{"role": "user", "content": "hi"}] + assert image_b64 is None + # ===================================================================== # _friendly_error — httpx transport failures @@ -488,13 +834,10 @@ class TestBuildPassthroughPayloadToolChoice: class TestFriendlyErrorHttpx: - """The async pass-through helpers talk to llama-server via httpx. - When the subprocess is down, httpx raises RequestError subclasses - whose string form (``"All connection attempts failed"``, ``"[Errno 111] - Connection refused"``, ...) does NOT contain the substring - ``"Lost connection to llama-server"`` the sync path uses, so the - previous substring-only `_friendly_error` returned a useless generic - message. These tests pin the new isinstance-based mapping. + """When llama-server is down, httpx RequestError strings lack the + "Lost connection to llama-server" substring the sync path keys off, so the + old substring-only `_friendly_error` returned a useless generic message. + These tests pin the new isinstance-based mapping. """ def _req(self): @@ -517,18 +860,13 @@ class TestFriendlyErrorHttpx: assert "Lost connection" in _friendly_error(exc) def test_non_httpx_unchanged(self): - # Non-httpx exceptions still fall through to the existing substring - # heuristics — a context-size message must still produce the - # "Message too long" path. - ctx_msg = ( - "request (4096 tokens) exceeds the available context size (2048 tokens)" - ) + # Non-httpx exceptions still fall through to the substring heuristics + # — a context-size message must still produce "Message too long". + ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)" assert "Message too long" in _friendly_error(ValueError(ctx_msg)) def test_generic_exception_returns_generic_message(self): - assert ( - _friendly_error(RuntimeError("unrelated")) == "An internal error occurred" - ) + assert _friendly_error(RuntimeError("unrelated")) == "An internal error occurred" from routes.inference import ( # noqa: E402 @@ -546,23 +884,17 @@ class TestDropEmptyAssistantSentinels: {"role": "user", "content": "again"}, ] out = _drop_empty_assistant_sentinels(msgs) - assert out == [ - {"role": "user", "content": "hi"}, - {"role": "user", "content": "again"}, - ] + assert out == [{"role": "user", "content": "hi"}, {"role": "user", "content": "again"}] def test_drops_assistant_with_no_content_key(self): - # exclude_none=True strips the content key entirely; filter must catch this. + # exclude_none=True strips the content key entirely; filter must catch it. msgs = [ {"role": "user", "content": "hi"}, {"role": "assistant"}, {"role": "user", "content": "ok"}, ] out = _drop_empty_assistant_sentinels(msgs) - assert out == [ - {"role": "user", "content": "hi"}, - {"role": "user", "content": "ok"}, - ] + assert out == [{"role": "user", "content": "hi"}, {"role": "user", "content": "ok"}] def test_preserves_assistant_with_text(self): msgs = [ @@ -662,22 +994,16 @@ class TestGgufVisionMessages: messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True) assert has_image is True - assert messages[0]["content"][0] == { - "type": "text", - "text": "describe image one", - } + assert messages[0]["content"][0] == {"type": "text", "text": "describe image one"} assert messages[0]["content"][1]["type"] == "image_url" assert len(messages[0]["content"]) == 2 - assert messages[2]["content"][0] == { - "type": "text", - "text": "describe image two", - } + assert messages[2]["content"][0] == {"type": "text", "text": "describe image two"} assert messages[2]["content"][1]["type"] == "image_url" assert len(messages[2]["content"]) == 2 assert isinstance(messages[1]["content"], str) - # Legacy top-level image_base64 must be ignored when any message-level - # image already exists; otherwise turn 2 ends up with two image parts. + # Legacy top-level image_base64 must be ignored when a message-level + # image exists; otherwise turn 2 ends up with two image parts. for msg in messages: content = msg.get("content") if isinstance(content, list): @@ -694,14 +1020,9 @@ class TestGgufVisionMessages: messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True) assert has_image is True - assert messages[0]["content"][0] == { - "type": "text", - "text": "describe this image", - } + assert messages[0]["content"][0] == {"type": "text", "text": "describe this image"} assert messages[0]["content"][1]["type"] == "image_url" - assert messages[0]["content"][1]["image_url"]["url"].startswith( - "data:image/png;base64," - ) + assert messages[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,") def test_rejects_image_parts_for_text_only_gguf(self): req = ChatCompletionRequest( @@ -725,3 +1046,252 @@ class TestGgufVisionMessages: with pytest.raises(HTTPException) as exc_info: _openai_messages_for_gguf_chat(req, is_vision = False) assert "does not support vision" in str(exc_info.value) + + def test_tool_nudge_system_update_preserves_image_parts(self): + messages = [ + {"role": "system", "content": "Base instructions."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{self._PNG_B64}", + }, + }, + ], + }, + ] + + updated = _set_or_prepend_system_message( + messages, "Base instructions.\n\nUse tools when appropriate." + ) + + assert updated[0] == { + "role": "system", + "content": "Base instructions.\n\nUse tools when appropriate.", + } + assert updated[1]["content"][1]["type"] == "image_url" + assert messages[1]["content"][1]["type"] == "image_url" + + def test_tool_nudge_system_update_handles_none_messages(self): + assert _set_or_prepend_system_message(None, "") == [] + assert _set_or_prepend_system_message(None, "Use tools.") == [ + {"role": "system", "content": "Use tools."} + ] + + def test_tool_nudge_system_update_dedupes_non_leading_system(self): + messages = [ + {"role": "user", "content": "earlier"}, + {"role": "system", "content": "Mid instructions."}, + {"role": "user", "content": "now"}, + ] + + updated = _set_or_prepend_system_message(messages, "Mid instructions.\n\nUse tools.") + + assert [m["role"] for m in updated] == ["system", "user", "user"] + assert updated[0]["content"] == "Mid instructions.\n\nUse tools." + + +class TestGgufVisionToolRouting: + class _Request: + async def is_disconnected(self): + return False + + @staticmethod + def _drive(coro): + return asyncio.run(coro) + + @staticmethod + def _consume_response(response): + async def _consume(): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + return chunks + + return TestGgufVisionToolRouting._drive(_consume()) + + def test_image_request_with_enabled_tools_enters_gguf_tool_loop(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + + def _plain(**kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**kwargs): + captured["kwargs"] = kwargs + yield {"type": "content", "text": "done"} + + backend = SimpleNamespace( + is_loaded = True, + is_vision = True, + supports_tools = True, + model_identifier = "gemma-4-12b-it-GGUF", + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + enable_tools = True, + enabled_tools = ["web_search"], + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + { + "type": "image_url", + "image_url": { + "url": (f"data:image/png;base64,{TestGgufVisionMessages._PNG_B64}"), + }, + }, + ], + }, + ], + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + self._consume_response(response) + + assert "kwargs" in captured + assert captured["kwargs"]["tools"] + tool_messages = captured["kwargs"]["messages"] + assert tool_messages[0]["role"] == "system" + assert tool_messages[1]["role"] == "user" + assert tool_messages[1]["content"][1]["type"] == "image_url" + + def test_parallel_tool_calls_false_reaches_gguf_tool_loop(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + + def _plain(**kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**kwargs): + captured["kwargs"] = kwargs + yield {"type": "content", "text": "done"} + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + model_identifier = "test-gguf", + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + enable_tools = True, + enabled_tools = ["web_search"], + parallel_tool_calls = False, + messages = [{"role": "user", "content": "search once"}], + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + self._consume_response(response) + + assert captured["kwargs"]["disable_parallel_tool_use"] is True + + def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): + import routes.inference as inf_mod + + captured = {} + + def _generate(**kwargs): + captured["messages"] = kwargs["messages"] + yield "done" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + "finish_reason": "stop", + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + model_identifier = "test-gguf", + generate_chat_completion = _generate, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [ + {"role": "system", "content": "original system"}, + {"role": "developer", "content": "developer rules"}, + {"role": "user", "content": "hi"}, + ], + ) + + self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + + assert captured["messages"] == [ + {"role": "system", "content": "original system\n\ndeveloper rules"}, + {"role": "user", "content": "hi"}, + ] + + @pytest.mark.parametrize( + ("seed", "expected"), + [ + (41, [41, 42, 43]), + (-1, [-1, -1, -1]), + ], + ) + def test_gguf_n_choices_vary_explicit_non_negative_seed(self, monkeypatch, seed, expected): + import routes.inference as inf_mod + + seen_seeds = [] + + def _generate(**kwargs): + seen_seeds.append(kwargs.get("seed")) + yield f"choice-{len(seen_seeds)}" + yield { + "type": "metadata", + "usage": { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12, + }, + "finish_reason": "stop", + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + model_identifier = "test-gguf", + generate_chat_completion = _generate, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + n = 3, + seed = seed, + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + body = json.loads(response.body) + + assert seen_seeds == expected + assert [choice["index"] for choice in body["choices"]] == [0, 1, 2] diff --git a/studio/backend/tests/test_openai_tool_result_fallbacks.py b/studio/backend/tests/test_openai_tool_result_fallbacks.py index 7c033bc348..5ea812441f 100644 --- a/studio/backend/tests/test_openai_tool_result_fallbacks.py +++ b/studio/backend/tests/test_openai_tool_result_fallbacks.py @@ -3,8 +3,8 @@ """Regression tests for OpenAI Responses tool-result rendering. -Covers two bug classes: empty web_search cards (per-card result seeded -with "Searching: ") and orphan shell_call cards (bundled-output +Two bug classes: empty web_search cards (per-card result seeded with +"Searching: ") and orphan shell_call cards (bundled-output fallback + final flush at response.completed / response.incomplete). """ @@ -137,8 +137,8 @@ def test_web_search_each_call_carries_its_own_query_as_result(monkeypatch): def test_web_search_last_call_overwritten_with_citations(monkeypatch): - """Last call still gets the aggregated citation list; earlier calls - keep their per-call `Searching:` text.""" + """Last call gets the aggregated citations; earlier calls keep their + per-call `Searching:` text.""" sse_events = [ { "type": "response.output_item.done", @@ -170,12 +170,12 @@ def test_web_search_last_call_overwritten_with_citations(monkeypatch): events = _tool_events(lines) ends = [e for e in events if e["type"] == "tool_end"] by_id: dict = {} - # Keep the LAST tool_end per id (the citation overwrite for ws_2). + # Keep the LAST tool_end per id (citation overwrite for ws_2). for e in ends: by_id[e["tool_call_id"]] = e # First call keeps its own query. assert by_id["ws_1"]["result"] == "Searching: first query" - # Last call gets overwritten with the citation block. + # Last call overwritten with the citation block. assert "Title: Example A" in by_id["ws_2"]["result"] assert "URL: https://example.com/a" in by_id["ws_2"]["result"] diff --git a/studio/backend/tests/test_pricing.py b/studio/backend/tests/test_pricing.py index 313c15a441..8cd7796f14 100644 --- a/studio/backend/tests/test_pricing.py +++ b/studio/backend/tests/test_pricing.py @@ -1,8 +1,8 @@ # 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 per-session cost calculator. Verifies math -against ``core/inference/pricing.py`` and graceful degradation.""" +"""Unit tests for the per-session cost calculator: math against +``core/inference/pricing.py`` plus graceful degradation.""" import math @@ -21,7 +21,11 @@ from core.inference.pricing import ( ) -def _isclose(a, b, tol = 1e-6): +def _isclose( + a, + b, + tol = 1e-6, +): return math.isclose(a, b, rel_tol = tol, abs_tol = tol) @@ -241,9 +245,7 @@ def test_openai_cache_read_subtracted_from_input_at_discount(): ) # 20k charged at full price, 80k charged at 0.1x assert _isclose(out["input_usd"], 20_000 / 1_000_000.0 * base) - assert _isclose( - out["cache_read_usd"], 80_000 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT - ) + assert _isclose(out["cache_read_usd"], 80_000 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT) def test_openai_billable_input_tokens_does_not_double_count_cache_read(): @@ -390,9 +392,7 @@ def test_openai_web_search_charged_per_thousand(): "openai_tool_use": {"web_search_requests": 250}, }, ) - assert _isclose( - out["server_tools_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K - ) + assert _isclose(out["server_tools_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K) assert _isclose(out["total_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K) @@ -426,8 +426,7 @@ def test_openai_tool_surcharges_added_to_total(): expected_input = 100_000 / 1_000_000.0 * 5.0 expected_output = 5_000 / 1_000_000.0 * 30.0 expected_tools = ( - 3 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K - + 0.25 * OPENAI_CONTAINER_USD_PER_HOUR + 3 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K + 0.25 * OPENAI_CONTAINER_USD_PER_HOUR ) assert _isclose( out["total_usd"], @@ -463,12 +462,12 @@ def test_snapshot_contains_provider_buckets_and_multipliers(): # ── longest-prefix match: dated mini variant must not collide with the -# shorter family prefix. ── +# shorter family prefix ── def test_longest_prefix_match_wins_for_dated_mini_snapshot(): - """`gpt-5.4-mini-2026-...` must inherit the mini rate, not the - shorter `gpt-5.4` rate (longest prefix wins).""" + """`gpt-5.4-mini-2026-...` inherits the mini rate, not the shorter + `gpt-5.4` rate (longest prefix wins).""" out = calculate_cost( "openai", "gpt-5.4-mini-2026-04-23", @@ -494,7 +493,7 @@ def test_longest_prefix_match_wins_for_dated_pro_snapshot(): def test_openai_chat_style_usage_keys_priced_correctly(): - """Chat-style envelope (`prompt_tokens` / `completion_tokens`) must + """Chat-style envelope (`prompt_tokens`/`completion_tokens`) must produce a non-zero cost (previously silently zeroed).""" out = calculate_cost( "openai", @@ -578,7 +577,7 @@ def test_openai_chat_style_prompt_tokens_keeps_cache_read_semantics(): def test_openai_chat_style_envelope_reads_cache_from_prompt_tokens_details(): """Chat-style envelope ships cached under prompt_tokens_details; - calculator must honour both this and input_tokens_details.""" + calculator must honour both that and input_tokens_details.""" base = OPENAI_PRICING["gpt-5.5"]["input_per_mtok"] raw = calculate_cost( "openai", @@ -600,10 +599,7 @@ def test_openai_chat_style_envelope_reads_cache_from_prompt_tokens_details(): ) # Both envelopes must price identically. assert _isclose(chat_style["input_usd"], raw["input_usd"]), (chat_style, raw) - assert _isclose(chat_style["cache_read_usd"], raw["cache_read_usd"]), ( - chat_style, - raw, - ) + assert _isclose(chat_style["cache_read_usd"], raw["cache_read_usd"]), (chat_style, raw) # 80k at 0.1x base, 20k at full. assert _isclose( chat_style["cache_read_usd"], diff --git a/studio/backend/tests/test_pricing_edge.py b/studio/backend/tests/test_pricing_edge.py index ca4be258e0..1fcc428f90 100644 --- a/studio/backend/tests/test_pricing_edge.py +++ b/studio/backend/tests/test_pricing_edge.py @@ -18,7 +18,11 @@ from core.inference.pricing import ( ) -def _isclose(a, b, tol = 1e-6): +def _isclose( + a, + b, + tol = 1e-6, +): return math.isclose(a, b, rel_tol = tol, abs_tol = tol) @@ -26,8 +30,8 @@ def _isclose(a, b, tol = 1e-6): def test_prefix_match_requires_dash_boundary_opus_variant(): - # `claude-opus-4-15` must not inherit `claude-opus-4-1` pricing; - # next char must be `-` or end-of-string. + # `claude-opus-4-15` must not inherit `claude-opus-4-1` pricing; the next + # char must be `-` or end-of-string. assert _lookup("anthropic", "claude-opus-4-15") is None out = calculate_cost( "anthropic", @@ -51,8 +55,8 @@ def test_prefix_match_requires_dash_boundary_gpt_variant(): def test_prefix_match_requires_dash_boundary_pro_lookalike(): - # `gpt-5.5-prod` must fall through `gpt-5.5-pro` (6x overcharge) - # and land on the canonical `gpt-5.5` row. + # `gpt-5.5-prod` must fall through `gpt-5.5-pro` (6x overcharge) and land on + # the canonical `gpt-5.5` row. prices = _lookup("openai", "gpt-5.5-prod") assert prices is not None assert ( @@ -77,7 +81,7 @@ def test_prefix_match_still_resolves_legit_dated_snapshots(): assert out["priced"] is True assert _isclose(out["input_usd"], 0.75) - # And Anthropic dated snapshot still resolves to canonical row. + # Anthropic dated snapshot still resolves to the canonical row. out = calculate_cost( "anthropic", "claude-opus-4-7-20260414", @@ -91,7 +95,6 @@ def test_prefix_match_still_resolves_legit_dated_snapshots(): def test_explicit_zero_input_tokens_wins_over_stale_prompt_tokens(): - # Input-side mirror of the output zero precedence test. out = calculate_cost( "openai", "gpt-5.5", @@ -106,7 +109,7 @@ def test_explicit_zero_input_tokens_wins_over_stale_prompt_tokens(): def test_none_input_tokens_falls_through_to_prompt_tokens(): - # `None` is "key present but unset"; chat-style mirror wins. + # `None` means "key present but unset"; chat-style mirror wins. out = calculate_cost( "openai", "gpt-5.5", @@ -172,8 +175,8 @@ def test_negative_prompt_tokens_chat_style_clamp(): def test_anthropic_chat_cache_read_exceeds_prompt_no_negative_billable(): - # cache_read > prompt_tokens clamps uncached_input at 0; billable - # still reflects cache buckets (we charge for what we got). + # cache_read > prompt_tokens clamps uncached_input at 0; billable still + # reflects cache buckets (we charge for what we got). out = calculate_cost( "anthropic", "claude-opus-4-7", @@ -188,9 +191,7 @@ def test_anthropic_chat_cache_read_exceeds_prompt_no_negative_billable(): assert out["billable_input_tokens"] == 500 # 0 uncached + 500 cache_read # cache_read still priced at the discount rate. base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"] - assert _isclose( - out["cache_read_usd"], 500 / 1_000_000.0 * base * ANTHROPIC_CACHE_READ_MULT - ) + assert _isclose(out["cache_read_usd"], 500 / 1_000_000.0 * base * ANTHROPIC_CACHE_READ_MULT) def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached(): @@ -207,17 +208,15 @@ def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached(): ) assert out["input_usd"] == 0.0 # Cache read still priced (the 0.1x bucket). - assert _isclose( - out["cache_read_usd"], 500 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT - ) + assert _isclose(out["cache_read_usd"], 500 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT) # ── long-context tier crosses on billable, including cache_creation ── def test_openai_long_context_triggers_on_cache_creation_inflated_billable(): - # cache_creation pushes billable past 272k -> long-context tier - # must fire to avoid undercounting. + # cache_creation pushes billable past 272k -> long-context tier must fire to + # avoid undercounting. out = calculate_cost( "openai", "gpt-5.5", @@ -272,8 +271,8 @@ def test_openai_chat_envelope_long_context_parity_with_raw(): def test_cache_creation_as_int_does_not_crash(): - # Proxies sometimes fold cache_creation to an int; tolerate it - # and fall back to the 5m default. + # Proxies sometimes fold cache_creation to an int; tolerate it and fall back + # to the 5m default. base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"] out = calculate_cost( "anthropic", @@ -363,16 +362,16 @@ def test_empty_usage_dict_zero_bill(): def test_anthropic_prompt_tokens_details_fallback_when_native_key_missing(): - """Chat-style envelope without `cache_read_input_tokens` but with - mirrored `prompt_tokens_details.cached_tokens` should still apply - the cache_read discount.""" + """Chat-style envelope without `cache_read_input_tokens` but with mirrored + `prompt_tokens_details.cached_tokens` should still apply the cache_read + discount.""" r = calculate_cost( provider = "anthropic", model = "claude-opus-4-7", usage = { "prompt_tokens": 1_000_000, "completion_tokens": 0, - # Only the mirrored shape (no native key). + # Mirrored shape only (no native key). "prompt_tokens_details": {"cached_tokens": 1_000_000}, "cache_creation_input_tokens": 0, }, @@ -383,8 +382,8 @@ def test_anthropic_prompt_tokens_details_fallback_when_native_key_missing(): def test_anthropic_native_key_takes_precedence_over_mirrored(): - """When both native and mirrored cache-read fields are present, - the native Anthropic field wins (mirror is fallback-only).""" + """When both native and mirrored cache-read fields are present, the native + Anthropic field wins (mirror is fallback-only).""" r = calculate_cost( provider = "anthropic", model = "claude-opus-4-7", @@ -403,8 +402,8 @@ def test_anthropic_native_key_takes_precedence_over_mirrored(): def test_anthropic_native_zero_takes_precedence_over_mirrored(): - """Explicit `cache_read_input_tokens: 0` is authoritative; a stale - mirrored block from a proxy must not inflate cache_read past it.""" + """Explicit `cache_read_input_tokens: 0` is authoritative; a stale mirrored + block from a proxy must not inflate cache_read past it.""" r = calculate_cost( provider = "anthropic", model = "claude-opus-4-7", @@ -429,8 +428,8 @@ def test_anthropic_native_zero_takes_precedence_over_mirrored(): def test_build_usage_chunk_forwards_anthropic_cache_creation_breakdown(): - """Chat-style envelope must carry the 5m/1h cache-write breakdown - so downstream cost calc applies the 2x 1h premium.""" + """Chat-style envelope must carry the 5m/1h cache-write breakdown so + downstream cost calc applies the 2x 1h premium.""" import json from core.inference.external_provider import _build_usage_chunk diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py index 0e668944f4..5e24ed752d 100644 --- a/studio/backend/tests/test_providers_api.py +++ b/studio/backend/tests/test_providers_api.py @@ -4,13 +4,13 @@ """ Integration tests for the external providers API. -Requires a running Unsloth Studio server. Configure via environment variables: +Requires a running Unsloth Studio server. Configure via env vars: export STUDIO_TEST_URL="http://localhost:8888" # default export STUDIO_TEST_USER="unsloth" # default export STUDIO_TEST_PASSWORD="..." # required — see .bootstrap_password - # Provider API keys — any left unset will have their tests automatically skipped + # Provider API keys — tests skip when their key is unset export OPENAI_API_KEY="sk-..." export MISTRAL_API_KEY="..." export GOOGLE_API_KEY="..." @@ -38,15 +38,14 @@ BASE_URL = os.getenv("STUDIO_TEST_URL", "http://localhost:8000") USERNAME = os.getenv("STUDIO_TEST_USER", "unsloth") PASSWORD = os.getenv("STUDIO_TEST_PASSWORD", "") -# These tests require a live Studio server reachable at BASE_URL with a known -# bootstrap password. Skip the whole module when that environment is missing -# (e.g. on CI runners) so pytest discovery does not error out. +# Skip the whole module when no live Studio server / bootstrap password is +# available (e.g. on CI) so pytest discovery does not error out. pytestmark = pytest.mark.skipif( not PASSWORD, reason = "Integration test requires a running Studio server; set STUDIO_TEST_PASSWORD to enable.", ) -# Map provider_type → (env var name, model to use for inference test) +# provider_type → (env var name, model for inference test) _PROVIDER_CONFIGS: dict[str, tuple[str, str]] = { "openai": ("OPENAI_API_KEY", "gpt-4o-mini"), "mistral": ("MISTRAL_API_KEY", "mistral-small-2506"), @@ -73,11 +72,9 @@ def _url(path: str) -> str: def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]: - """ - Read a streaming SSE response and return (assembled_text, saw_done). + """Read an SSE response, return (assembled_text, saw_done). - Each chunk is a JSON object with choices[0].delta.content. - The stream ends with `data: [DONE]`. + Each chunk is JSON with choices[0].delta.content; stream ends at `data: [DONE]`. """ reply_parts: list[str] = [] saw_done = False @@ -93,7 +90,7 @@ def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]: break try: chunk = json.loads(data) - # Handle both error payloads and normal chunks + # Handle error payloads and normal chunks if "error" in chunk: raise RuntimeError(f"Provider error in stream: {chunk['error']}") delta = chunk.get("choices", [{}])[0].get("delta", {}) @@ -111,20 +108,12 @@ def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]: @pytest.fixture(scope = "session") def auth_headers() -> dict[str, str]: - """ - Log in once per session and return auth headers. + """Log in once per session and return auth headers. - On a fresh Studio install the bootstrap password triggers a forced password - change (must_change_password=True). Any subsequent API call using that token - returns 403 "Password change required". This fixture detects that state, - automatically completes the change-password flow, and re-logs in so all other - tests get a fully usable token. - - The new password used during auto-change is: - STUDIO_TEST_NEW_PASSWORD (env var, optional) - or PASSWORD + "-test" (derived default) - - On the second run, set STUDIO_TEST_PASSWORD to the new password. + On a fresh install the bootstrap password forces a change; this fixture + detects must_change_password, auto-completes the change (new password = + STUDIO_TEST_NEW_PASSWORD or PASSWORD + "-test"), and re-logs in. On the + second run, set STUDIO_TEST_PASSWORD to the new password. """ assert PASSWORD, ( "STUDIO_TEST_PASSWORD is not set.\n" @@ -142,8 +131,8 @@ def auth_headers() -> dict[str, str]: assert token, "access_token is empty" if body.get("must_change_password"): - # Bootstrap token is restricted — only /api/auth/change-password works with it. - # Auto-complete the forced change so the rest of the tests get a full token. + # Bootstrap token only works with change-password; auto-complete the + # forced change so the rest of the tests get a full token. new_password = os.getenv("STUDIO_TEST_NEW_PASSWORD") or f"{PASSWORD}-test" change_resp = requests.post( _url("/api/auth/change-password"), @@ -175,12 +164,10 @@ def public_key_pem(auth_headers: dict[str, str]) -> str: @pytest.fixture(scope = "session") def vision_image_data_url() -> str: - """ - Download the sloth image once per session and return it as a base64 data URI. + """Download the sloth image once per session as a base64 data URI. - Using a data URI instead of a remote URL ensures every provider receives - the image inline — Gemini's OpenAI-compatible layer does not fetch external - HTTP URLs, so raw image_url links silently produce empty replies for Gemini. + A data URI sends the image inline; Gemini's OpenAI-compatible layer doesn't + fetch external HTTP URLs, so raw image_url links give empty Gemini replies. """ resp = requests.get(_VISION_IMAGE_URL, timeout = 30) resp.raise_for_status() @@ -191,11 +178,10 @@ def vision_image_data_url() -> str: @pytest.fixture(scope = "session") def encrypt_key(public_key_pem: str): + """Return encrypt_key(plaintext) -> base64 RSA-OAEP ciphertext. + + Uses the backend's RSA public key; mirrors the frontend. """ - Return a callable encrypt_key(plaintext: str) -> str (base64 RSA-OAEP ciphertext). - Uses the backend's RSA public key — mirrors what the frontend does. - """ - # Decode PEM → load RSA public key pem_bytes = public_key_pem.encode("utf-8") rsa_pub = serialization.load_pem_public_key(pem_bytes) @@ -225,9 +211,7 @@ class TestAuth: json = {"username": USERNAME, "password": PASSWORD}, timeout = 10, ) - assert ( - resp.status_code == 200 - ), f"Login failed ({resp.status_code}): {resp.text}" + assert resp.status_code == 200, f"Login failed ({resp.status_code}): {resp.text}" body = resp.json() assert body.get("access_token"), "access_token is missing or empty" assert body.get("token_type") == "bearer" @@ -237,9 +221,7 @@ class TestAuth: class TestPublicKey: - def test_public_key_is_valid_pem( - self, auth_headers: dict[str, str], public_key_pem: str - ): + def test_public_key_is_valid_pem(self, auth_headers: dict[str, str], public_key_pem: str): """GET /api/providers/public-key returns an importable RSA PEM key.""" pem_bytes = public_key_pem.encode("utf-8") key = serialization.load_pem_public_key(pem_bytes) @@ -261,9 +243,7 @@ class TestRegistry: ) assert resp.status_code == 200, f"Registry failed: {resp.text}" providers = resp.json() - assert ( - len(providers) == 9 - ), f"Expected 9 providers, got {len(providers)}: {providers}" + assert len(providers) == 9, f"Expected 9 providers, got {len(providers)}: {providers}" print(f"\n {'Provider':<12} {'Base URL'}") print(f" {'-'*12} {'-'*45}") for p in providers: @@ -283,9 +263,7 @@ class TestRegistry: def test_registry_entries_have_required_fields(self, auth_headers: dict[str, str]): """Each registry entry has provider_type, display_name, base_url, default_models.""" - resp = requests.get( - _url("/api/providers/registry"), headers = auth_headers, timeout = 10 - ) + resp = requests.get(_url("/api/providers/registry"), headers = auth_headers, timeout = 10) assert resp.status_code == 200 for entry in resp.json(): for field in ( @@ -306,8 +284,8 @@ class TestRegistry: class TestProviderCRUD: """ - These tests run sequentially within the class and share state via class variables. - They create, read, update, and delete a single test provider config. + Run sequentially, sharing state via class variables. Create, read, update, + and delete a single test provider config. """ _created_id: str = "" @@ -320,9 +298,7 @@ class TestProviderCRUD: json = {"provider_type": "openai", "display_name": "Test OpenAI (pytest)"}, timeout = 10, ) - assert ( - resp.status_code == 201 - ), f"Create failed ({resp.status_code}): {resp.text}" + assert resp.status_code == 201, f"Create failed ({resp.status_code}): {resp.text}" body = resp.json() assert body.get("id"), "No id in response" assert body["provider_type"] == "openai" @@ -333,9 +309,7 @@ class TestProviderCRUD: def test_list_includes_created(self, auth_headers: dict[str, str]): """GET /api/providers/ includes the newly created config.""" - assert ( - TestProviderCRUD._created_id - ), "No created_id (run test_create_provider first)" + assert TestProviderCRUD._created_id, "No created_id (run test_create_provider first)" resp = requests.get(_url("/api/providers/"), headers = auth_headers, timeout = 10) assert resp.status_code == 200 ids = [p["id"] for p in resp.json()] @@ -354,9 +328,7 @@ class TestProviderCRUD: json = {"display_name": new_name}, timeout = 10, ) - assert ( - resp.status_code == 200 - ), f"Update failed ({resp.status_code}): {resp.text}" + assert resp.status_code == 200, f"Update failed ({resp.status_code}): {resp.text}" assert resp.json()["display_name"] == new_name print(f"\n updated display_name to '{new_name}'") @@ -368,14 +340,10 @@ class TestProviderCRUD: headers = auth_headers, timeout = 10, ) - assert ( - resp.status_code == 204 - ), f"Delete failed ({resp.status_code}): {resp.text}" + assert resp.status_code == 204, f"Delete failed ({resp.status_code}): {resp.text}" - # Confirm gone from list - list_resp = requests.get( - _url("/api/providers/"), headers = auth_headers, timeout = 10 - ) + # Confirm gone + list_resp = requests.get(_url("/api/providers/"), headers = auth_headers, timeout = 10) ids = [p["id"] for p in list_resp.json()] assert TestProviderCRUD._created_id not in ids, "Deleted provider still in list" print(f"\n deleted id={TestProviderCRUD._created_id} confirmed gone") @@ -384,7 +352,7 @@ class TestProviderCRUD: # ── TestProviderInference ──────────────────────────────────────────── -# Build parametrize list: (provider_type, model, api_key) for configured providers only +# Parametrize (provider_type, model, api_key) for configured providers _INFERENCE_PARAMS = [ pytest.param( ptype, @@ -402,8 +370,8 @@ _INFERENCE_PARAMS = [ class TestProviderInference: """ - Live inference tests — one parametrized set per provider. - Each test is automatically skipped when the provider's API key env var is not set. + Live inference tests, one parametrized set per provider. Each is skipped + when the provider's API key env var is unset. """ @pytest.mark.parametrize("provider_type,model,api_key", _INFERENCE_PARAMS) @@ -423,9 +391,7 @@ class TestProviderInference: json = {"provider_type": provider_type, "encrypted_api_key": encrypted}, timeout = 30, ) - assert ( - resp.status_code == 200 - ), f"Request failed ({resp.status_code}): {resp.text}" + assert resp.status_code == 200, f"Request failed ({resp.status_code}): {resp.text}" body = resp.json() assert ( body["success"] is True @@ -449,9 +415,7 @@ class TestProviderInference: json = {"provider_type": provider_type, "encrypted_api_key": encrypted}, timeout = 30, ) - assert ( - resp.status_code == 200 - ), f"Request failed ({resp.status_code}): {resp.text}" + assert resp.status_code == 200, f"Request failed ({resp.status_code}): {resp.text}" models = resp.json() assert isinstance(models, list), f"Expected list, got {type(models)}" assert len(models) > 0, f"No models returned for {provider_type}" @@ -497,10 +461,8 @@ class TestProviderInference: # ── TestVisionInference ───────────────────────────────────────────── -# Sloth photo — used to test vision routing across providers -_VISION_IMAGE_URL = ( - "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg" -) +# Sloth photo for testing vision routing across providers +_VISION_IMAGE_URL = "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg" _VISION_PARAMS = [ pytest.param( @@ -520,8 +482,8 @@ _VISION_PARAMS = [ class TestVisionInference: """ - Send a 1×1 white PNG alongside a text question to each vision-capable provider. - Verifies that image content parts survive the proxy and the provider replies. + Send a 1×1 white PNG plus a text question to each vision-capable provider. + Verifies image content parts survive the proxy and the provider replies. """ @pytest.mark.parametrize("provider_type,model,api_key", _VISION_PARAMS) @@ -580,12 +542,10 @@ class TestVisionInference: class TestLocalInferenceUnaffected: def test_chat_without_provider(self, auth_headers: dict[str, str]): - """ - POST /v1/chat/completions without provider fields must not return 422 or 500. + """POST /v1/chat/completions without provider fields must not 422/500. - 200 = a local model is loaded and responded. - 503 = no model loaded (expected in test environment — that's fine). - Any other 4xx/5xx (except 503) = regression in request handling. + 200 = local model responded; 503 = no model loaded (fine in tests); + any other 4xx/5xx = request-handling regression. """ resp = requests.post( _url("/v1/chat/completions"), @@ -602,8 +562,6 @@ class TestLocalInferenceUnaffected: f"This likely means the provider fields broke the base request schema." ) status_label = ( - "local model responded" - if resp.status_code == 200 - else "no model loaded (expected)" + "local model responded" if resp.status_code == 200 else "no model loaded (expected)" ) print(f"\n status={resp.status_code} ({status_label}) — local path unaffected") diff --git a/studio/backend/tests/test_pytorch_mirror.py b/studio/backend/tests/test_pytorch_mirror.py index 5844f209b6..59842214ef 100644 --- a/studio/backend/tests/test_pytorch_mirror.py +++ b/studio/backend/tests/test_pytorch_mirror.py @@ -19,8 +19,8 @@ OFFICIAL_URL = "https://download.pytorch.org/whl" def _reload_whl_base(monkeypatch, mirror_value = None): - """(Re-)import install_python_stack with a controlled env and return _PYTORCH_WHL_BASE.""" - # Remove cached module so the module-level assignment re-executes + """(Re-)import install_python_stack with a controlled env, return _PYTORCH_WHL_BASE.""" + # Drop cached module so the module-level assignment re-executes. sys.modules.pop("install_python_stack", None) if mirror_value is None: @@ -28,7 +28,7 @@ def _reload_whl_base(monkeypatch, mirror_value = None): else: monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", mirror_value) - # Temporarily add the script's directory to sys.path for import + # Add the script's directory to sys.path for import. script_dir = str(_INSTALL_SCRIPT.parent) monkeypatch.syspath_prepend(script_dir) diff --git a/studio/backend/tests/test_rag_captioning.py b/studio/backend/tests/test_rag_captioning.py new file mode 100644 index 0000000000..5d83a7d38d --- /dev/null +++ b/studio/backend/tests/test_rag_captioning.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Multimodal captioning tests: gating, grouping, splice, retrieval.""" + +from __future__ import annotations + +from core.rag import captioner +from core.rag.parsers import Page, ParsedImage + + +def _img(page): + return ParsedImage(image_bytes = b"\x89PNG fake", page_number = page, xref = page) + + +def test_caption_images_disabled_by_default(monkeypatch): + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) + assert captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) == {} + + +def test_caption_images_groups_by_page(monkeypatch): + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 8) + monkeypatch.setattr(captioner, "_caption_one", lambda base, model, b, t: "a chart of results") + out = captioner.caption_images([_img(1), _img(1), _img(3)], endpoint = ("http://x", "local")) + assert out == {1: ["a chart of results", "a chart of results"], 3: ["a chart of results"]} + + +def test_caption_images_respects_cap(monkeypatch): + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 2) + calls = [] + monkeypatch.setattr(captioner, "_caption_one", lambda *a: (calls.append(1) or "cap")) + captioner.caption_images([_img(i) for i in range(5)], endpoint = ("http://x", "local")) + assert len(calls) == 2 + + +def test_caption_images_no_endpoint(monkeypatch): + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) + assert captioner.caption_images([_img(1)]) == {} + + +def test_splice_captions_appends_to_right_page(): + pages = [Page("body one", 1, 8), Page("body two", 2, 8)] + out = captioner.splice_captions(pages, {2: ["a diagram of X"]}) + assert out[0].text == "body one" + assert "a diagram of X" in out[1].text + assert out[1].text.startswith("body two") + assert out[1].char_count == len(out[1].text) + + +def test_splice_captions_noop_when_empty(): + pages = [Page("body", 1, 4)] + assert captioner.splice_captions(pages, {}) is pages + + +def test_render_pdf_figures_detects_drawing(tmp_path): + import pymupdf + + from core.rag.parsers import render_pdf_figures + + pdf = tmp_path / "fig.pdf" + doc = pymupdf.open() + page = doc.new_page() + shape = page.new_shape() + shape.draw_rect(pymupdf.Rect(60, 60, 540, 460)) + for i in range(8): + shape.draw_line((80, 80 + i * 40), (520, 80 + i * 40)) + shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) + shape.commit() + doc.save(str(pdf)) + doc.close() + + figs = render_pdf_figures(str(pdf)) + assert figs, "expected at least one rendered figure region" + assert figs[0].image_bytes[:8] == b"\x89PNG\r\n\x1a\n" + assert figs[0].page_number == 1 + + +def test_captioned_text_is_searchable(rag_home, stub_embeddings, monkeypatch): + from core.rag import retrieval, store + from storage import rag_db + + pages = [Page("Section 1 intro text about models.", 1, 33)] + pages = captioner.splice_captions( + pages, {1: ["bar chart comparing throughput across quantizations"]} + ) + from core.rag import chunking, embeddings + + chunks = chunking.chunk_pages( + pages, max_tokens = 128, overlap = 16, count = embeddings.token_counter(None) + ) + vecs = embeddings.encode([c.text for c in chunks], normalize = True) + + conn = rag_db.get_connection() + try: + kb_id = store.create_kb(conn, name = "kb") + scope = store.kb_scope(kb_id) + doc_id = store.create_document(conn, scope = scope, filename = "d.pdf", sha256 = "h") + store.add_chunks(conn, scope, doc_id, chunks, vecs) + hits = retrieval.retrieve_lexical(conn, scope, "throughput quantizations", k = 5) + finally: + conn.close() + assert hits, "spliced caption text should be retrievable via lexical search" diff --git a/studio/backend/tests/test_rag_chunking.py b/studio/backend/tests/test_rag_chunking.py new file mode 100644 index 0000000000..3d3c9eedd8 --- /dev/null +++ b/studio/backend/tests/test_rag_chunking.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Chunking unit tests (no DB, no model).""" + +from core.rag.chunking import chunk_pages +from core.rag.parsers import Page, parse_text + +WORDS = lambda t: len(t.split()) # noqa: E731 + + +def _page(text: str, page_number = None) -> Page: + return Page(text = text, page_number = page_number, char_count = len(text)) + + +def test_chunk_token_bounds_and_overlap(): + text = " ".join(f"w{i}" for i in range(300)) + chunks = chunk_pages([_page(text)], max_tokens = 128, overlap = 24, count = WORDS) + assert len(chunks) >= 3 + assert all(c.token_count <= 128 for c in chunks) + a, b = chunks[0].text.split(), chunks[1].text.split() + shared = next((n for n in range(60, 0, -1) if a[-n:] == b[:n]), 0) + assert shared == 24 # exactly overlap tokens carried + + +def test_chunk_never_exceeds_max_with_overlap_carry(): + """Overlap carry is trimmed so no chunk exceeds max_tokens (else the embedder overflows).""" + s1 = " ".join("a" for _ in range(10)) + s2 = " ".join("b" for _ in range(95)) # near max + chunks = chunk_pages([_page(f"{s1}. {s2}")], max_tokens = 100, overlap = 24, count = WORDS) + assert all(c.token_count <= 100 for c in chunks), [c.token_count for c in chunks] + + +def test_chunk_indices_are_sequential(): + chunks = chunk_pages([_page("alpha. " * 200)], max_tokens = 32, overlap = 0, count = WORDS) + assert [c.chunk_index for c in chunks] == list(range(len(chunks))) + + +def test_chunk_tracks_source_page_index(): + pages = [_page("alpha bravo " * 80, 1), _page("charlie delta " * 80, 2)] + chunks = chunk_pages(pages, max_tokens = 32, overlap = 0, count = WORDS) + page0 = [c for c in chunks if c.source_page_index == 0] + page1 = [c for c in chunks if c.source_page_index == 1] + assert page0 and page1 + assert all(c.page_number == 1 for c in page0) + assert all(c.page_number == 2 for c in page1) + + +def test_chunk_char_offsets_locate_text_in_page(): + # Each chunk's char span must slice back to text containing it. + page_text = "alpha bravo charlie delta echo foxtrot golf hotel " * 30 + pages = [_page(page_text, 1)] + chunks = chunk_pages(pages, max_tokens = 16, overlap = 0, count = WORDS) + assert len(chunks) > 1 + for c in chunks: + assert 0 <= c.page_char_start < c.page_char_end <= len(page_text) + sliced = page_text[c.page_char_start : c.page_char_end] + assert c.text in sliced or sliced.strip() == c.text + + +def test_empty_page_yields_no_chunks(): + chunks = chunk_pages([_page(" \n ")], max_tokens = 32, overlap = 0, count = WORDS) + assert chunks == [] + + +def test_parse_text_single_page(): + pages = parse_text("hello world") + assert len(pages) == 1 + assert pages[0].char_count == len("hello world") diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py new file mode 100644 index 0000000000..8321068afd --- /dev/null +++ b/studio/backend/tests/test_rag_embed_llama_server.py @@ -0,0 +1,448 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""llama-server GGUF embedder tests, every boundary mocked.""" + +import subprocess +import sys +import textwrap +from pathlib import Path + +import numpy as np +import pytest + +from core.rag import config, embeddings +from core.rag import embed_llama_server as mod +from core.rag.embed_llama_server import LlamaServerBackend + + +@pytest.fixture(autouse = True) +def _reset_backend_singleton(): + embeddings._reset_backend() + yield + embeddings._reset_backend() + + +class _FakeProc: + """subprocess.Popen stand-in with controllable liveness.""" + + def __init__( + self, + alive = True, + returncode = 0, + ): + self._alive = alive + self.returncode = returncode + self.stdout = iter(()) # drain thread exits immediately + + def poll(self): + return None if self._alive else self.returncode + + def terminate(self): + self._alive = False + + def kill(self): + self._alive = False + + def wait(self, timeout = None): + return self.returncode + + +def _mock_auto(monkeypatch, *, gpus, binary): + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setattr(config, "EMBED_BACKEND", "auto") + monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: gpus)) + monkeypatch.setattr(LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary)) + + +def _stub_st_load(monkeypatch): + # Make the ST probe succeed without importing sentence-transformers (absent in + # the torch-free backend CI); these tests assert selection, not a real load. + monkeypatch.setattr(embeddings, "_get", lambda *a, **k: object()) + + +def test_auto_uses_st_with_cuda(monkeypatch): + _stub_st_load(monkeypatch) + _mock_auto(monkeypatch, gpus = [(0, 40000)], binary = "/bin/llama-server") + assert type(embeddings._get_backend()).__name__ == "_SentenceTransformersBackend" + + +def test_auto_uses_llama_without_cuda(monkeypatch): + _mock_auto(monkeypatch, gpus = [], binary = "/bin/llama-server") + assert isinstance(embeddings._get_backend(), LlamaServerBackend) + + +def test_auto_falls_back_to_st_without_binary(monkeypatch): + _stub_st_load(monkeypatch) + _mock_auto(monkeypatch, gpus = [], binary = None) + assert type(embeddings._get_backend()).__name__ == "_SentenceTransformersBackend" + + +def test_llama_backend_selected_by_config(monkeypatch): + monkeypatch.setattr(config, "EMBED_BACKEND", "llama-server") + assert isinstance(embeddings._get_backend(), LlamaServerBackend) + + +def test_unknown_backend_raises(monkeypatch): + monkeypatch.setattr(config, "EMBED_BACKEND", "bogus") + with pytest.raises(ValueError, match = "Unknown RAG_EMBED_BACKEND"): + embeddings._get_backend() + + +def test_explicit_backend_overrides_auto(monkeypatch): + _stub_st_load(monkeypatch) + monkeypatch.setattr(config, "EMBED_BACKEND", "sentence-transformers") + assert type(embeddings._get_backend()).__name__ == "_SentenceTransformersBackend" + monkeypatch.setattr(config, "EMBED_BACKEND", "llama-server") + assert isinstance(embeddings._get_backend(), LlamaServerBackend) + + +def test_llama_backend_imports_no_torch(): + # Clean subprocess so the parent's imports don't mask a regression. + backend_dir = Path(__file__).resolve().parents[1] + code = textwrap.dedent( + """ + import sys + from core.rag import embeddings + b = embeddings._get_backend() + assert type(b).__name__ == "LlamaServerBackend", type(b).__name__ + assert "torch" not in sys.modules, "torch was imported" + assert "sentence_transformers" not in sys.modules, "ST was imported" + print("OK") + """ + ) + env = { + **__import__("os").environ, + "RAG_EMBED_BACKEND": "llama-server", + "PYTHONPATH": str(backend_dir), + } + proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True, env = env) + assert proc.returncode == 0, proc.stderr + assert "OK" in proc.stdout + + +def test_build_cmd_cpu_flags(): + b = LlamaServerBackend() + cmd = b._build_cmd("/bin/llama-server", "/m/bge.gguf", 9999, use_gpu = False) + assert "--embedding" in cmd + assert cmd[cmd.index("--pooling") + 1] == "cls" + assert cmd[cmd.index("--fit") + 1] == "off" # deterministic, no auto-resize + assert cmd[cmd.index("-ngl") + 1] == "0" # CPU keeps all off the GPU + assert cmd[cmd.index("--port") + 1] == "9999" + + +def test_build_cmd_gpu_offloads(): + b = LlamaServerBackend() + cmd = b._build_cmd("/bin/llama-server", "/m/bge.gguf", 1, use_gpu = True) + assert cmd[cmd.index("-ngl") + 1] == "-1" # offload all, matching the chat server + + +def test_build_env_cpu_hides_gpus(): + b = LlamaServerBackend() + env = b._build_env("/bin/llama-server", use_gpu = False) + assert env["CUDA_VISIBLE_DEVICES"] == "" # never contend with the chat model + assert env["LLAMA_SET_ROWS"] == "1" + + +def test_build_env_gpu_inherits_devices(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + b = LlamaServerBackend() + env = b._build_env("/bin/llama-server", use_gpu = True) + assert env.get("CUDA_VISIBLE_DEVICES") == "0,1" # inherit Studio's selection + + +def test_use_gpu_explicit_modes(monkeypatch): + b = LlamaServerBackend() + monkeypatch.setattr(config, "EMBED_DEVICE", "gpu") + assert b._use_gpu() is True + monkeypatch.setattr(config, "EMBED_DEVICE", "cpu") + assert b._use_gpu() is False + + +def test_use_gpu_auto_follows_probe(monkeypatch): + b = LlamaServerBackend() + monkeypatch.setattr(config, "EMBED_DEVICE", "auto") + monkeypatch.setattr(LlamaServerBackend, "_gpu_available", staticmethod(lambda: True)) + assert b._use_gpu() is True + monkeypatch.setattr(LlamaServerBackend, "_gpu_available", staticmethod(lambda: False)) + assert b._use_gpu() is False + + +def test_use_gpu_sticky_cpu_fallback(monkeypatch): + b = LlamaServerBackend() + monkeypatch.setattr(config, "EMBED_DEVICE", "auto") + monkeypatch.setattr(LlamaServerBackend, "_gpu_available", staticmethod(lambda: True)) + b._force_cpu = True # a prior GPU start failed + assert b._use_gpu() is False + + +def test_gpu_available_reuses_studio_probe(monkeypatch): + import utils.hardware as uh + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setattr(uh, "is_apple_silicon", lambda: False) + # Ample free VRAM -> GPU; nearly full -> CPU; none -> CPU. + monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: [(0, 40000)])) + assert LlamaServerBackend._gpu_available() is True + monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: [(0, 100)])) + assert LlamaServerBackend._gpu_available() is False + monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: [])) + assert LlamaServerBackend._gpu_available() is False + + +def test_gpu_available_apple_metal(monkeypatch): + import utils.hardware as uh + monkeypatch.setattr(uh, "is_apple_silicon", lambda: True) + assert LlamaServerBackend._gpu_available() is True + + +def _patch_spawn_deps( + monkeypatch, + proc, + *, + free_port = 54321, +): + # Force CPU so spawn never depends on a host GPU. + monkeypatch.setattr(config, "EMBED_DEVICE", "cpu") + monkeypatch.setattr(LlamaServerBackend, "_resolve_binary", lambda self: "/bin/llama-server") + monkeypatch.setattr(LlamaServerBackend, "_resolve_model_path", lambda self: "/m/bge.gguf") + monkeypatch.setattr(LlamaServerBackend, "_find_free_port", staticmethod(lambda: free_port)) + monkeypatch.setattr(mod.subprocess, "Popen", lambda *a, **k: proc) + + +def test_spawn_uses_explicit_port(monkeypatch): + monkeypatch.setattr(config, "EMBED_PORT", 8123) + b = LlamaServerBackend() + _patch_spawn_deps(monkeypatch, _FakeProc(alive = True)) + monkeypatch.setattr(b, "_wait_for_health", lambda *a, **k: True) + b._spawn() + assert b._port == 8123 + + +def test_spawn_uses_free_port_when_auto(monkeypatch): + monkeypatch.setattr(config, "EMBED_PORT", 0) + b = LlamaServerBackend() + _patch_spawn_deps(monkeypatch, _FakeProc(alive = True), free_port = 47000) + monkeypatch.setattr(b, "_wait_for_health", lambda *a, **k: True) + b._spawn() + assert b._port == 47000 + + +def test_spawn_fails_loud_on_early_exit(monkeypatch): + monkeypatch.setattr(config, "EMBED_PORT", 8124) + b = LlamaServerBackend() + _patch_spawn_deps(monkeypatch, _FakeProc(alive = False, returncode = 1)) + with pytest.raises(RuntimeError, match = "failed to become healthy"): + b._spawn() + + +def test_spawn_auto_falls_back_to_cpu_on_gpu_failure(monkeypatch): + monkeypatch.setattr(config, "EMBED_DEVICE", "auto") + monkeypatch.setattr(LlamaServerBackend, "_gpu_available", staticmethod(lambda: True)) + b = LlamaServerBackend() + calls = [] + + def fake_spawn_once(use_gpu): + calls.append(use_gpu) + if use_gpu: + raise RuntimeError("CUDA out of memory") + + monkeypatch.setattr(b, "_spawn_once", fake_spawn_once) + b._spawn() + assert calls == [True, False] # tried GPU, then fell back to CPU + assert b._force_cpu is True # sticky, so respawns stay on CPU + + +def test_spawn_explicit_gpu_does_not_fall_back(monkeypatch): + monkeypatch.setattr(config, "EMBED_DEVICE", "gpu") + b = LlamaServerBackend() + + def fake_spawn_once(use_gpu): + raise RuntimeError("CUDA out of memory") + + monkeypatch.setattr(b, "_spawn_once", fake_spawn_once) + with pytest.raises(RuntimeError, match = "out of memory"): + b._spawn() + assert b._force_cpu is False # explicit gpu never silently downgrades + + +def _embed_response(vectors): + # Reversed so the index sort is exercised. + items = [{"index": i, "embedding": v} for i, v in enumerate(vectors)] + return {"data": list(reversed(items))} + + +def test_encode_orders_and_returns_float32(monkeypatch): + b = LlamaServerBackend() + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + captured = {} + + def fake_post(path, payload): + captured["path"] = path + captured["input"] = payload["input"] + return _embed_response([[3.0, 4.0], [0.0, 5.0]]) + + monkeypatch.setattr(b, "_post", fake_post) + out = b.encode(["a", "b"], normalize = False) + assert captured["path"] == "/v1/embeddings" + assert out.dtype == np.float32 + assert out.shape == (2, 2) + assert out[0].tolist() == [3.0, 4.0] # index sort restored order + + +def test_encode_normalizes(monkeypatch): + b = LlamaServerBackend() + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + monkeypatch.setattr(b, "_post", lambda p, pl: _embed_response([[3.0, 4.0]])) + out = b.encode(["a"], normalize = True) + np.testing.assert_allclose(np.linalg.norm(out, axis = 1), [1.0], rtol = 1e-6) + + +def test_encode_empty_returns_zero_rows(monkeypatch): + b = LlamaServerBackend() + b._dim = 384 + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + out = b.encode([]) + assert out.shape == (0, 384) + assert out.dtype == np.float32 + + +def test_encode_rejects_count_mismatch(monkeypatch): + b = LlamaServerBackend() + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + monkeypatch.setattr(b, "_post", lambda p, pl: {"data": [{"index": 0, "embedding": [1.0]}]}) + with pytest.raises(RuntimeError, match = "vectors for"): + b.encode(["a", "b"], normalize = False) + + +def test_encode_batches(monkeypatch): + monkeypatch.setattr(config, "EMBED_BATCH", 2) + b = LlamaServerBackend() + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + calls = [] + + def fake_post(path, payload): + chunk = payload["input"] + calls.append(len(chunk)) + return _embed_response([[1.0, 0.0]] * len(chunk)) + + monkeypatch.setattr(b, "_post", fake_post) + out = b.encode(["a", "b", "c"], normalize = False) + assert out.shape == (3, 2) + assert calls == [2, 1] # batched at EMBED_BATCH=2 + + +def test_dim_probes_once_and_caches(monkeypatch): + b = LlamaServerBackend() + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + n_calls = {"n": 0} + + def fake_post(path, payload): + n_calls["n"] += 1 + return _embed_response([[0.1] * 384]) + + monkeypatch.setattr(b, "_post", fake_post) + assert b.dim() == 384 + assert b.dim() == 384 + assert n_calls["n"] == 1 # cached after the first probe + + +def test_token_counter_hits_tokenize(monkeypatch): + b = LlamaServerBackend() + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + seen = {} + + def fake_post(path, payload): + seen["path"] = path + seen["content"] = payload["content"] + return {"tokens": [1, 2, 3, 4]} + + monkeypatch.setattr(b, "_post", fake_post) + count = b.token_counter() + assert count("hello world") == 4 + assert seen["path"] == "/tokenize" + assert seen["content"] == "hello world" + + +def test_ensure_ready_respawns_dead_process(monkeypatch): + b = LlamaServerBackend() + b._process = _FakeProc(alive = False, returncode = 0) + spawned = {"n": 0} + + def fake_spawn(): + spawned["n"] += 1 + b._process = _FakeProc(alive = True) + + monkeypatch.setattr(b, "_spawn", fake_spawn) + b._ensure_ready() + assert spawned["n"] == 1 + assert b._process_alive() + # Already alive -> no second spawn. + b._ensure_ready() + assert spawned["n"] == 1 + + +def test_post_restarts_once_on_connect_error(monkeypatch): + import httpx + + b = LlamaServerBackend() + b._port = 9000 + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + restarts = {"n": 0} + monkeypatch.setattr(b, "_restart", lambda: restarts.__setitem__("n", restarts["n"] + 1)) + + attempts = {"n": 0} + + class _Client: + def post(self, url, json): + attempts["n"] += 1 + if attempts["n"] == 1: + raise httpx.ConnectError("boom") + + class _R: + def raise_for_status(self_inner): + return None + + def json(self_inner): + return {"tokens": [1]} + + return _R() + + b._client = _Client() + out = b._post("/tokenize", {"content": "x"}) + assert out == {"tokens": [1]} + assert restarts["n"] == 1 # one self-heal restart, then success + + +def test_post_restarts_once_on_read_timeout(monkeypatch): + # A wedged request (ReadTimeout) also triggers one restart-and-retry. + import httpx + + b = LlamaServerBackend() + b._port = 9000 + monkeypatch.setattr(b, "_ensure_ready", lambda: None) + restarts = {"n": 0} + monkeypatch.setattr(b, "_restart", lambda: restarts.__setitem__("n", restarts["n"] + 1)) + + attempts = {"n": 0} + + class _Client: + def post(self, url, json): + attempts["n"] += 1 + if attempts["n"] == 1: + raise httpx.ReadTimeout("timed out") + + class _R: + def raise_for_status(self_inner): + return None + + def json(self_inner): + return {"data": [{"index": 0, "embedding": [1.0, 0.0]}]} + + return _R() + + b._client = _Client() + out = b._post("/v1/embeddings", {"input": ["x"]}) + assert out["data"][0]["embedding"] == [1.0, 0.0] + assert restarts["n"] == 1 # timeout self-heals like a transport error diff --git a/studio/backend/tests/test_rag_embeddings.py b/studio/backend/tests/test_rag_embeddings.py new file mode 100644 index 0000000000..28a2f69426 --- /dev/null +++ b/studio/backend/tests/test_rag_embeddings.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Embedder concurrency tests: the fast tokenizer isn't thread-safe, so encode +and token counting must be serialized (else threads panic "Already borrowed").""" + +import os +import threading +import time + +import numpy as np +import pytest + +from core.rag import config, embeddings + + +@pytest.fixture(autouse = True) +def _pin_st_backend(monkeypatch): + # Tests patch ST internals (_get), so force the ST backend. + monkeypatch.setattr(config, "EMBED_BACKEND", "sentence-transformers") + embeddings._reset_backend() + yield + embeddings._reset_backend() + + +class _ConcurrencyProbe: + """Records whether two callers were in the guarded body at once.""" + + def __init__(self): + self.inside = 0 + self.saw_overlap = False + self._g = threading.Lock() + + def enter(self): + with self._g: + self.inside += 1 + if self.inside > 1: + self.saw_overlap = True + time.sleep(0.005) # widen the race window + with self._g: + self.inside -= 1 + + +class _FakeModel: + def __init__(self, probe): + self._probe = probe + self.tokenizer = _FakeTokenizer(probe) + + def encode(self, texts, **_kw): + self._probe.enter() + return np.zeros((len(texts), 4), dtype = np.float32) + + +class _FakeTokenizer: + def __init__(self, probe): + self._probe = probe + + def encode(self, text, **_kw): + self._probe.enter() + return list(range(len(text.split()))) + + +def _hammer(fn, n = 8): + errors: list[Exception] = [] + + def worker(): + try: + fn() + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target = worker) for _ in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + return errors + + +def test_encode_is_serialized(monkeypatch): + probe = _ConcurrencyProbe() + monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _FakeModel(probe)) + errors = _hammer(lambda: embeddings.encode(["alpha beta", "gamma"])) + assert errors == [] + assert probe.saw_overlap is False # compute lock serialized encode() + + +def test_token_counter_is_serialized(monkeypatch): + probe = _ConcurrencyProbe() + monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _FakeModel(probe)) + count = embeddings.token_counter() + errors = _hammer(lambda: count("one two three four")) + assert errors == [] + assert probe.saw_overlap is False # counting shares the tokenizer lock + + +def test_encode_enables_parallelism_only_during_call(monkeypatch): + seen = {} + + class _M: + tokenizer = None + + def encode(self, texts, **_kw): + seen["during"] = os.environ.get("TOKENIZERS_PARALLELISM") + return np.zeros((len(texts), 4), dtype = np.float32) + + monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _M()) + os.environ["TOKENIZERS_PARALLELISM"] = "false" + embeddings.encode(["alpha", "beta"]) + assert seen["during"] == "true" # rayon batch tokenization enabled in-call + assert os.environ.get("TOKENIZERS_PARALLELISM") == "false" # restored after + + +def test_token_counter_enables_parallelism_only_during_call(monkeypatch): + seen = {} + + class _Tok: + def encode(self, text, **_kw): + seen["during"] = os.environ.get("TOKENIZERS_PARALLELISM") + return list(range(len(text.split()))) + + class _M: + tokenizer = _Tok() + + monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _M()) + os.environ["TOKENIZERS_PARALLELISM"] = "false" + count = embeddings.token_counter() + count("alpha beta gamma") + assert seen["during"] == "true" # rayon enabled in-call, like _st_encode + assert os.environ.get("TOKENIZERS_PARALLELISM") == "false" # restored after + + +class _SentinelLlamaBackend: + """Stand-in for LlamaServerBackend; never spawns a real server.""" + + +def _force_st_load_failure(monkeypatch): + """Make the ST warm-probe raise.""" + + def _boom(model_name = None): + raise RuntimeError("torch is broken on this machine") + + monkeypatch.setattr(embeddings, "_get", _boom) + + +def _patch_llama_backend(monkeypatch, *, binary): + from core.inference.llama_cpp import LlamaCppBackend + from core.rag import embed_llama_server + + monkeypatch.setattr(LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary)) + monkeypatch.setattr(embed_llama_server, "LlamaServerBackend", _SentinelLlamaBackend) + + +def test_st_failure_falls_back_to_llama_server(monkeypatch): + # ST can't load but llama-server is available -> use it. + _force_st_load_failure(monkeypatch) + _patch_llama_backend(monkeypatch, binary = "/fake/llama-server") + embeddings._reset_backend() + backend = embeddings._get_backend() + assert isinstance(backend, _SentinelLlamaBackend) + + +def test_st_failure_without_llama_binary_reraises(monkeypatch): + # No llama-server binary -> surface the failure, don't degrade to nothing. + _force_st_load_failure(monkeypatch) + _patch_llama_backend(monkeypatch, binary = None) + embeddings._reset_backend() + with pytest.raises(RuntimeError, match = "torch is broken"): + embeddings._get_backend() + + +def test_st_success_keeps_sentence_transformers(monkeypatch): + # Clean ST probe -> ST backend stays selected, no fallback. + monkeypatch.setattr(embeddings, "_get", lambda model_name = None: object()) + _patch_llama_backend(monkeypatch, binary = "/fake/llama-server") + embeddings._reset_backend() + backend = embeddings._get_backend() + assert isinstance(backend, embeddings._SentenceTransformersBackend) + + +class _BoomOnEncodeModel: + """Loads fine (init probe passes) but raises when encoding.""" + + tokenizer = None + + def encode(self, texts, **_kw): + raise RuntimeError("CUDA error during encode") + + +def test_st_encode_runtime_failure_switches_to_llama(monkeypatch): + # encode() blows up mid-run -> switch to llama-server and stay switched. + monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _BoomOnEncodeModel()) + _patch_llama_backend(monkeypatch, binary = "/fake/llama-server") + calls = {} + + def _sentinel_encode( + self, + texts, + *, + model_name = None, + normalize = True, + ): + calls["used"] = True + return np.zeros((len(texts), 4), dtype = np.float32) + + monkeypatch.setattr(_SentinelLlamaBackend, "encode", _sentinel_encode, raising = False) + embeddings._reset_backend() + + out = embeddings.encode(["alpha", "beta"]) + assert calls.get("used") is True # retried on the llama fallback + assert out.shape == (2, 4) + # Switch is process-wide: later calls keep using llama, not ST. + assert isinstance(embeddings._get_backend(), _SentinelLlamaBackend) + + +def test_st_encode_failure_without_llama_binary_reraises(monkeypatch): + # No llama-server binary -> surface the encode error. + monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _BoomOnEncodeModel()) + _patch_llama_backend(monkeypatch, binary = None) + embeddings._reset_backend() + with pytest.raises(RuntimeError, match = "CUDA error during encode"): + embeddings.encode(["alpha", "beta"]) diff --git a/studio/backend/tests/test_rag_ingestion.py b/studio/backend/tests/test_rag_ingestion.py new file mode 100644 index 0000000000..c5fe15876d --- /dev/null +++ b/studio/backend/tests/test_rag_ingestion.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Ingestion lifecycle tests: pending -> completed, SSE events, dedupe, delete.""" + +import os +import time + +import pytest + +from core.rag import ingestion, store +from storage import rag_db + + +def _write(tmp_path, name, text): + path = tmp_path / name + path.write_text(text, encoding = "utf-8") + return str(path) + + +def _drain(job_id): + return list(ingestion.job_events(job_id)) + + +def _wait_completed(job_id, timeout = 30.0): + deadline = time.time() + timeout + while time.time() < deadline: + status = ingestion.get_job_status(job_id) + if status and status["status"] in ("completed", "failed"): + return status + time.sleep(0.05) + raise AssertionError("ingestion did not finish in time") + + +def test_ingestion_lifecycle_pending_to_completed(rag_home, stub_embeddings, tmp_path): + path = _write(tmp_path, "doc.txt", "alpha bravo charlie " * 50) + scope = store.kb_scope("K1") + doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + + conn = rag_db.get_connection() + try: + assert store.get_document(conn, doc_id)["status"] == "pending" + finally: + conn.close() + + events = _drain(job_id) + assert any(e["type"] == "progress" for e in events) + assert events[-1]["type"] == "complete" + assert events[-1]["num_chunks"] > 0 + + status = _wait_completed(job_id) + assert status["status"] == "completed" + assert status["progress"] == 1.0 + + conn = rag_db.get_connection() + try: + doc = store.get_document(conn, doc_id) + assert doc["status"] == "completed" + assert doc["num_chunks"] > 0 + assert store.search_lexical(conn, scope, "alpha", 10) + finally: + conn.close() + + +def test_ingestion_dedupe_by_hash(rag_home, stub_embeddings, tmp_path): + path = _write(tmp_path, "doc.txt", "alpha bravo charlie") + scope = store.kb_scope("K1") + doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + _drain(job_id) + _wait_completed(job_id) + + # Identical content -> same doc id, no re-ingest. + path2 = _write(tmp_path, "copy.txt", "alpha bravo charlie") + doc_id2, job_id2 = ingestion.start_ingestion(scope, "K1", None, "copy.txt", path2) + events = _drain(job_id2) + assert doc_id2 == doc_id + assert any(e.get("deduped") for e in events) + + conn = rag_db.get_connection() + try: + assert len(store.list_documents(conn, scope)) == 1 + finally: + conn.close() + + +def test_ingestion_delete_removes_all_rows(rag_home, stub_embeddings, tmp_path): + path = _write(tmp_path, "doc.txt", "alpha bravo charlie delta") + scope = store.kb_scope("K1") + doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + _drain(job_id) + _wait_completed(job_id) + + conn = rag_db.get_connection() + try: + store.delete_document(conn, doc_id) + assert store.get_document(conn, doc_id) is None + assert store.search_lexical(conn, scope, "alpha", 10) == [] + assert store.list_documents(conn, scope) == [] + finally: + conn.close() + + +def test_ingestion_rejects_unsupported_ext(rag_home, stub_embeddings, tmp_path): + path = _write(tmp_path, "doc.xyz", "alpha") + with pytest.raises(ValueError): + ingestion.start_ingestion(store.kb_scope("K1"), "K1", None, "doc.xyz", path) + + +def test_ingestion_empty_doc_completes_with_zero_chunks(rag_home, stub_embeddings, tmp_path): + path = _write(tmp_path, "empty.txt", " \n ") + scope = store.kb_scope("K1") + doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "empty.txt", path) + events = _drain(job_id) + assert events[-1]["type"] == "complete" + assert events[-1]["num_chunks"] == 0 + status = _wait_completed(job_id) + assert status["status"] == "completed" + + +@pytest.mark.skipif( + os.environ.get("RAG_REAL_EMBEDDER") != "1", + reason = "set RAG_REAL_EMBEDDER=1 to run the real sentence-transformers test", +) +def test_ingestion_with_real_embedder(rag_home, tmp_path): + path = _write(tmp_path, "doc.txt", "The Kestrel-9 turbine is rated at 9.5 megawatts.") + scope = store.kb_scope("K1") + doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + _drain(job_id) + status = _wait_completed(job_id, timeout = 120.0) + assert status["status"] == "completed" + + from core.rag import retrieval + + conn = rag_db.get_connection() + try: + hits = retrieval.retrieve_hybrid(conn, scope, "how much power does the turbine make?", k = 5) + assert hits and hits[0].chunk_id == f"{doc_id}:0" + finally: + conn.close() diff --git a/studio/backend/tests/test_rag_preview.py b/studio/backend/tests/test_rag_preview.py new file mode 100644 index 0000000000..e7f2a39792 --- /dev/null +++ b/studio/backend/tests/test_rag_preview.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""PDF region locator + citation preview route tests.""" + +from __future__ import annotations + +import time + +import pytest + +pytest.importorskip("pymupdf") +pytest.importorskip("sqlite_vec") + + +def _make_pdf(path) -> None: + import pymupdf + + doc = pymupdf.open() + body = ( + "BERT is designed to pre-train deep bidirectional representations.\n" + "The two pre-training objectives are masked language modeling and next " + "sentence prediction.\n" + "The Transformer base model uses eight attention heads in each layer.\n" + ) + for _ in range(3): + page = doc.new_page() + page.insert_text((72, 72), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def _ingest(home, pdf_path): + from core.rag import ingestion, store + from storage import rag_db + + conn = rag_db.get_connection() + kb_id = store.create_kb(conn, name = "kb") + conn.close() + doc_id, job_id = ingestion.start_ingestion( + store.kb_scope(kb_id), kb_id, None, "doc.pdf", str(pdf_path) + ) + t0 = time.time() + while time.time() - t0 < 30: + s = ingestion.get_job_status(job_id) + if s and s["status"] in ("completed", "failed"): + break + time.sleep(0.05) + assert s and s["status"] == "completed", s + return kb_id, doc_id + + +def test_chunks_carry_pdf_regions(rag_home, stub_embeddings): + from utils.paths import ensure_dir, rag_uploads_root + + pdf = ensure_dir(rag_uploads_root()) / "doc.pdf" + _make_pdf(pdf) + kb_id, doc_id = _ingest(rag_home, pdf) + + from storage import rag_db + + conn = rag_db.get_connection() + try: + rows = conn.execute( + "SELECT pdf_regions_json FROM chunks WHERE document_id=?", (doc_id,) + ).fetchall() + stored_path = conn.execute( + "SELECT stored_path FROM documents WHERE id=?", (doc_id,) + ).fetchone()["stored_path"] + finally: + conn.close() + + assert rows, "no chunks were stored" + assert stored_path and stored_path.endswith(".pdf") + with_regions = [r for r in rows if r["pdf_regions_json"]] + assert with_regions, "expected at least one chunk with PDF highlight regions" + import json + + region = json.loads(with_regions[0]["pdf_regions_json"])[0] + for key in ("pageIndex", "x", "y", "width", "height"): + assert key in region + if key in ("x", "y", "width", "height"): + assert 0.0 <= region[key] <= 1.0 + + +def test_preview_routes_and_signed_file(rag_home, stub_embeddings): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from auth.authentication import get_current_subject + from routes.rag import router + + from utils.paths import ensure_dir, rag_uploads_root + + pdf = ensure_dir(rag_uploads_root()) / "doc.pdf" + _make_pdf(pdf) + kb_id, doc_id = _ingest(rag_home, pdf) + + app = FastAPI() + app.include_router(router, prefix = "/api/rag") + app.dependency_overrides[get_current_subject] = lambda: "tester" + c = TestClient(app) + + res = c.post( + "/api/rag/search", + json = { + "query": "masked language modeling next sentence", + "kb_id": kb_id, + "mode": "lexical", + }, + ).json()["results"] + assert res + chunk_id = res[0]["chunkId"] + + pt = c.get(f"/api/rag/documents/{doc_id}/preview-target", params = {"chunk_id": chunk_id}).json() + assert pt["mediaKind"] == "pdf" + assert pt["text"] + + url = c.get(f"/api/rag/documents/{doc_id}/file-url").json()["url"] + full = c.get(url) + assert full.status_code == 200 and full.content[:4] == b"%PDF" + rng = c.get(url, headers = {"Range": "bytes=0-99"}) + assert rng.status_code in (200, 206) + assert ( + c.get( + f"/api/rag/documents/{doc_id}/file-signed", + params = {"token": "bad.token.sig"}, + ).status_code + == 401 + ) + + +def test_norm_token_decomposes_ligatures(): + # NFKC folds ligature glyphs to ASCII so anchors match (search_for misses these). + from core.rag.locators import _norm_token + + assert _norm_token("significant") == "significant" # fi + assert _norm_token("effort.") == "effort" # ff + trailing punct + assert _norm_token("**Bold**") == "bold" + assert _norm_token("...") == "" + + +def test_locator_handles_midword_anchor_and_locates_line(): + # A span beginning mid-word still locates: first/last tokens dropped. + import pymupdf + + from core.rag.locators import LocatorMatch, _regions_for_match + + doc = pymupdf.open() + page = doc.new_page() + page.insert_text((72, 200), "alpha beta gamma delta epsilon zeta eta theta", fontsize = 12) + page_text = doc[0].get_text("text") # mirrors what the parser stores + start = page_text.index("lpha") + end = page_text.index("theta") + 3 + match = LocatorMatch(page_index = 0, page_number = 1, start = start, end = end) + rects = _regions_for_match(doc, page_text, match) + doc.close() + + assert rects, "expected a located region for the interior phrase" + r = rects[0] + for k in ("pageIndex", "pageNumber", "x", "y", "width", "height"): + assert k in r + # Drawn near y=200 on a ~842pt page -> normalized y in the top half. + assert 0.0 < r["y"] < 0.5 + assert r["width"] > 0 and r["height"] > 0 + + +def test_sign_verify_roundtrip(rag_home): + from routes import rag as rag_routes + + tok = rag_routes._sign_document("doc-123") + assert rag_routes._verify_document_token(tok) == "doc-123" + assert rag_routes._verify_document_token("doc-123.0.deadbeef") is None # expired/bad + assert rag_routes._verify_document_token("garbage") is None diff --git a/studio/backend/tests/test_rag_retrieval.py b/studio/backend/tests/test_rag_retrieval.py new file mode 100644 index 0000000000..69d9e90871 --- /dev/null +++ b/studio/backend/tests/test_rag_retrieval.py @@ -0,0 +1,432 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Retrieval + tool tests: RRF fusion, min-score floor, scope, source-map.""" + +import math + +import pytest + +from core.rag import config, retrieval, store, tool +from core.rag.chunking import Chunk + +VOCAB = ["alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel"] + + +def _embed(text): + v = [float(text.lower().count(w)) for w in VOCAB] + n = math.sqrt(sum(x * x for x in v)) or 1.0 + return [x / n for x in v] + + +@pytest.fixture +def bow_embeddings(monkeypatch): + """Bag-of-words embedder matching the vectors stored in the db.""" + from core.rag import embeddings + + monkeypatch.setattr( + embeddings, + "encode", + lambda texts, *, model_name = None, normalize = True: [_embed(t) for t in texts], + ) + monkeypatch.setattr(embeddings, "dim", lambda model_name = None: len(VOCAB)) + + +def _chunk( + text, + index = 0, + page = None, +): + return Chunk( + text = text, + token_count = len(text.split()), + page_number = page, + source_page_index = 0, + chunk_index = index, + page_char_start = 0, + page_char_end = len(text), + ) + + +def _add_doc( + conn, + scope, + doc_id, + filename, + sha, + text, + page = None, +): + store.create_document(conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id) + store.add_chunks(conn, scope, doc_id, [_chunk(text, 0, page)], [_embed(text)]) + + +def test_rrf_ranks_doc_in_both_lists_first(): + # A chunk near the top of both rankings beats one in a single list. + lexical = [ + retrieval.Hit("a", 1.0, lexical_score = 1.0), + retrieval.Hit("b", 0.5, lexical_score = 0.5), + ] + dense = [ + retrieval.Hit("a", 0.9, dense_score = 0.9), + retrieval.Hit("c", 0.8, dense_score = 0.8), + ] + fused = retrieval._rrf([lexical, dense], rrf_k = 60, top_k = 10) + assert fused[0].chunk_id == "a" + assert fused[0].lexical_score == 1.0 and fused[0].dense_score == 0.9 + + +def test_retrieve_hybrid_returns_relevant_chunk(rag_conn, bow_embeddings): + _add_doc(rag_conn, "kb_a", "d1", "f1", "h1", "alpha bravo charlie") + _add_doc(rag_conn, "kb_a", "d2", "f2", "h2", "golf hotel delta") + hits = retrieval.retrieve_hybrid(rag_conn, "kb_a", "alpha bravo", k = 5) + assert hits[0].chunk_id == "d1:0" + + +def test_retrieve_dense_round_trips(rag_conn, bow_embeddings): + _add_doc(rag_conn, "kb_a", "d1", "f", "h1", "alpha alpha") + _add_doc(rag_conn, "kb_a", "d2", "f", "h2", "hotel golf") + hits = retrieval.retrieve_dense(rag_conn, "kb_a", "alpha", 5) + assert hits[0].chunk_id == "d1:0" + assert hits[0].dense_score is not None and hits[0].dense_score > 0.99 + + +def test_filter_min_score_gates_dense_hits(): + hits = [ + retrieval.Hit("a", 1.0, dense_score = 0.9), + retrieval.Hit("b", 0.5, dense_score = 0.2), + retrieval.Hit("c", 0.4, lexical_score = 0.4), # no dense_score -> kept + ] + out = retrieval.filter_min_score(hits, 0.5) + ids = {h.chunk_id for h in out} + assert ids == {"a", "c"} # b below floor, c lexical-only passes + assert retrieval.filter_min_score(hits, 0.0) == hits # floor off = identity + + +def test_tool_kb_scope_wins_over_thread(rag_conn, bow_embeddings, monkeypatch): + seen = {} + + def fake(conn, scope, q, **k): + seen["scope"] = scope + return [] + + monkeypatch.setattr(retrieval, "retrieve_hybrid", fake) + tool.search_knowledge_base(query = "q", scope_kb_id = "K", scope_thread_id = "T") + assert seen["scope"] == "kb_K" + + +def test_tool_empty_query_errors(rag_home): + assert tool.search_knowledge_base(query = " ").startswith("Error") + + +def test_tool_missing_scope_message(rag_home): + out = tool.search_knowledge_base(query = "hello") + assert "No documents" in out + + +def test_tool_formats_chunks_and_sources(rag_conn, bow_embeddings, monkeypatch): + _add_doc(rag_conn, "kb_a", "d1", "paper.pdf", "h1", "body text here", page = 3) + monkeypatch.setattr( + retrieval, + "retrieve_hybrid", + lambda conn, scope, q, **k: [retrieval.Hit("d1:0", 1.0)], + ) + text, sources = tool.search_knowledge_base_with_sources(query = "q", scope_kb_id = "a") + assert '' in text + assert "body text here" in text + assert sources == [ + { + "citationId": 1, + "chunkId": "d1:0", + "documentId": "d1", + "filename": "paper.pdf", + "page": 3, + "text": "body text here", + "score": 1.0, + } + ] + + +def test_tool_kb_scope_retrieves_from_db(rag_conn, bow_embeddings): + # End-to-end (no retrieve stub): doc found via its scope_kb_id (#8). + _add_doc(rag_conn, "kb_K", "d1", "kb.pdf", "h1", "alpha bravo charlie", page = 1) + text, sources = tool.search_knowledge_base_with_sources(query = "alpha bravo", scope_kb_id = "K") + assert "No matching chunks" not in text + assert sources and sources[0]["chunkId"] == "d1:0" + assert sources[0]["filename"] == "kb.pdf" + # A different KB id sees nothing (scope isolation). + other, other_sources = tool.search_knowledge_base_with_sources( + query = "alpha bravo", scope_kb_id = "OTHER" + ) + assert other_sources == [] and "No matching chunks" in other + + +def test_dispatcher_appends_sources_sentinel(rag_conn, bow_embeddings, monkeypatch): + # JSON source-map appended after the sentinel; text before it stays clean. + import json + + from core.inference import tools + + _add_doc(rag_conn, "kb_a", "d1", "paper.pdf", "h1", "body text here", page = 3) + monkeypatch.setattr( + retrieval, + "retrieve_hybrid", + lambda conn, scope, q, **k: [retrieval.Hit("d1:0", 1.0)], + ) + out = tools._search_knowledge_base({"query": "q"}, {"kb_id": "a"}) + assert tools.RAG_SOURCES_SENTINEL in out + model_text, _, payload = out.partition(tools.RAG_SOURCES_SENTINEL) + assert "__RAG_SOURCES__" not in model_text # model never sees the JSON + assert ' injected. + monkeypatch.setattr(retrieval, "retrieve_hybrid", _hits(0.8, key = "dense_score")) + found = tool.search_for_autoinject(query = "q", scope_kb_id = "a", min_dense_score = 0.55) + assert found is not None + text, sources = found + assert ' nothing injected. + monkeypatch.setattr(retrieval, "retrieve_hybrid", _hits(0.30, key = "dense_score")) + assert tool.search_for_autoinject(query = "q", scope_kb_id = "a", min_dense_score = 0.55) is None + + # Lexical-only hit (no dense score) does not auto-inject. + monkeypatch.setattr(retrieval, "retrieve_hybrid", _hits(1.0, key = "lexical_score")) + assert tool.search_for_autoinject(query = "q", scope_kb_id = "a", min_dense_score = 0.55) is None + + +def test_search_for_autoinject_bm25_gates_on_dense_probe(rag_conn, bow_embeddings, monkeypatch): + # BM25 hits carry no cosine, so the gate uses a dense 1-NN probe (#5). + _add_doc(rag_conn, "kb_a", "d1", "paper.pdf", "h1", "body text here", page = 3) + monkeypatch.setattr( + retrieval, + "retrieve_hybrid", + lambda conn, scope, q, **k: [retrieval.Hit("d1:0", 1.0, lexical_score = 2.5)], + ) + + monkeypatch.setattr( + retrieval, + "retrieve_dense", + lambda conn, scope, q, k = None, **kw: [retrieval.Hit("d1:0", 0.82, dense_score = 0.82)], + ) + found = tool.search_for_autoinject( + query = "q", scope_kb_id = "a", mode = "lexical", min_dense_score = 0.70 + ) + assert found is not None and found[1][0]["chunkId"] == "d1:0" + + monkeypatch.setattr( + retrieval, + "retrieve_dense", + lambda conn, scope, q, k = None, **kw: [retrieval.Hit("d1:0", 0.40, dense_score = 0.40)], + ) + assert ( + tool.search_for_autoinject(query = "q", scope_kb_id = "a", mode = "lexical", min_dense_score = 0.70) + is None + ) + + +def test_search_for_autoinject_empty_query_or_scope(rag_home): + assert tool.search_for_autoinject(query = " ", scope_kb_id = "a") is None + assert tool.search_for_autoinject(query = "hello") is None # no scope + + +def test_build_rag_autoinject_emits_pipeline(monkeypatch): + # Auto-inject yields the same tool card + source-map a real call would. + from core.inference import tools + from storage import rag_db + + monkeypatch.setattr(rag_db, "RAG_AVAILABLE", True, raising = False) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **k: ( + 'hi', + [{"citationId": 1, "filename": "d.pdf"}], + ), + ) + conv = [{"role": "user", "content": "When was DeepSeek V4 released?"}] + out = tools.build_rag_autoinject(conv, {"thread_id": "t1"}) + assert out is not None + kinds = [e["type"] for e in out["events"]] + assert "tool_start" in kinds and "tool_end" in kinds + te = next(e for e in out["events"] if e["type"] == "tool_end") + assert te["tool_name"] == "search_knowledge_base" + assert tools.RAG_SOURCES_SENTINEL in te["result"] + assert out["messages"][0]["tool_calls"][0]["function"]["name"] == "search_knowledge_base" + assert "__RAG_SOURCES__" not in out["messages"][1]["content"] + + +def test_build_rag_autoinject_skips_without_hit(monkeypatch): + from core.inference import tools + from storage import rag_db + + monkeypatch.setattr(rag_db, "RAG_AVAILABLE", True, raising = False) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **k: None) + assert ( + tools.build_rag_autoinject([{"role": "user", "content": "hi"}], {"thread_id": "t1"}) is None + ) + + +def test_build_rag_autoinject_enabled_by_default(monkeypatch): + from core.inference import tools + from storage import rag_db + + monkeypatch.delenv("RAG_AUTOINJECT", raising = False) + monkeypatch.delenv("RAG_AUTOINJECT_MIN_SCORE", raising = False) + monkeypatch.setattr(rag_db, "RAG_AVAILABLE", True, raising = False) + seen: dict = {} + + def fake(**k): + seen.update(k) + return ("x", [{"citationId": 1}]) + + monkeypatch.setattr(tool, "search_for_autoinject", fake) + out = tools.build_rag_autoinject([{"role": "user", "content": "hi"}], {"thread_id": "t1"}) + assert out is not None + assert seen["min_dense_score"] == 0.70 # high-precision floor by default + + +def test_build_rag_autoinject_caps_top_k(monkeypatch): + from core.inference import tools + from storage import rag_db + + monkeypatch.setenv("RAG_AUTOINJECT", "1") + monkeypatch.setenv("RAG_AUTOINJECT_TOP_K", "4") + monkeypatch.setattr(rag_db, "RAG_AVAILABLE", True, raising = False) + seen: dict = {} + + def fake(**k): + seen.update(k) + return ("x", [{"citationId": 1}]) + + monkeypatch.setattr(tool, "search_for_autoinject", fake) + conv = [{"role": "user", "content": "q"}] + tools.build_rag_autoinject(conv, {"thread_id": "t1"}) + assert seen["top_k"] == 4 # lean default + tools.build_rag_autoinject(conv, {"thread_id": "t1", "default_top_k": 2}) + assert seen["top_k"] == 2 # lower user setting wins + + +def test_build_rag_autoinject_disabled_by_env(monkeypatch): + from core.inference import tools + + monkeypatch.setenv("RAG_AUTOINJECT", "0") + assert ( + tools.build_rag_autoinject([{"role": "user", "content": "hi"}], {"thread_id": "t1"}) is None + ) + # No scope -> also a no-op. + monkeypatch.delenv("RAG_AUTOINJECT", raising = False) + assert tools.build_rag_autoinject([{"role": "user", "content": "hi"}], None) is None + + +def test_retrieve_hybrid_mode_selects_backend(monkeypatch): + # ``mode`` runs only the chosen backend; hybrid uses config counts + rrf_k. + calls: list = [] + monkeypatch.setattr( + retrieval, + "retrieve_lexical", + lambda c, s, q, k = None: calls.append(("lex", k)) or [], + ) + monkeypatch.setattr( + retrieval, + "retrieve_dense", + lambda c, s, q, k = None, *, model_name = None: calls.append(("dense", k)) or [], + ) + monkeypatch.setattr( + retrieval, + "_rrf", + lambda rankings, rrf_k, top_k: calls.append(("rrf", rrf_k, top_k)) or [], + ) + + calls.clear() + retrieval.retrieve_hybrid(None, "kb_a", "q", k = 5, mode = "lexical") + assert [c[0] for c in calls] == ["lex"] # dense + rrf skipped + + calls.clear() + retrieval.retrieve_hybrid(None, "kb_a", "q", k = 5, mode = "dense") + assert [c[0] for c in calls] == ["dense"] + + calls.clear() + retrieval.retrieve_hybrid(None, "kb_a", "q", k = 5, mode = "hybrid") + # Candidate pools + rrf_k come from config (no per-request override). + assert ("lex", config.TOP_K_LEXICAL) in calls + assert ("dense", config.TOP_K_DENSE) in calls + rrf = next(c for c in calls if c[0] == "rrf") + assert rrf[1] == config.RRF_K and rrf[2] == 5 # config rrf_k + final top_k + + +def test_scope_overrides_reach_retrieval(monkeypatch): + from core.inference import tools + from storage import rag_db + + monkeypatch.setattr(rag_db, "RAG_AVAILABLE", True, raising = False) + seen: dict = {} + + def fake_search(**kw): + seen.update(kw) + return ("text", []) + + monkeypatch.setattr(tool, "search_knowledge_base_with_sources", fake_search) + tools._search_knowledge_base( + {"query": "q"}, + {"kb_id": "a", "mode": "dense", "default_top_k": 11}, + ) + assert seen["mode"] == "dense" + assert seen["top_k"] == 11 + # Unknown mode falls back to hybrid. + seen.clear() + tools._search_knowledge_base({"query": "q"}, {"kb_id": "a", "mode": "bogus"}) + assert seen["mode"] == "hybrid" + + +def test_build_rag_autoinject_scope_overrides_env(monkeypatch): + from core.inference import tools + from storage import rag_db + + monkeypatch.setattr(rag_db, "RAG_AVAILABLE", True, raising = False) + seen: dict = {} + + def fake_autoinject(**k): + seen.update(k) + return ('hi', [{"citationId": 1}]) + + monkeypatch.setattr(tool, "search_for_autoinject", fake_autoinject) + conv = [{"role": "user", "content": "q"}] + + # Scope enables + overrides the floor though env says off. + monkeypatch.setenv("RAG_AUTOINJECT", "0") + out = tools.build_rag_autoinject( + conv, + { + "thread_id": "t1", + "autoinject": True, + "autoinject_min_score": 0.8, + "mode": "dense", + }, + ) + assert out is not None + assert seen["min_dense_score"] == 0.8 + assert seen["mode"] == "dense" + + # Explicit False disables even with the env default on. + monkeypatch.setenv("RAG_AUTOINJECT", "1") + assert tools.build_rag_autoinject(conv, {"thread_id": "t1", "autoinject": False}) is None diff --git a/studio/backend/tests/test_rag_store.py b/studio/backend/tests/test_rag_store.py new file mode 100644 index 0000000000..4c54b02ea7 --- /dev/null +++ b/studio/backend/tests/test_rag_store.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Store tests: incremental writes, dedupe, delete, scope, dense + lexical.""" + +import math + +from core.rag import store +from core.rag.chunking import Chunk + +VOCAB = ["alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel"] + + +def embed(text): + v = [float(text.lower().count(w)) for w in VOCAB] + n = math.sqrt(sum(x * x for x in v)) or 1.0 + return [x / n for x in v] + + +def _chunk( + text, + index = 0, + page = None, +): + return Chunk( + text = text, + token_count = len(text.split()), + page_number = page, + source_page_index = 0, + chunk_index = index, + page_char_start = 0, + page_char_end = len(text), + ) + + +def _add_doc(conn, scope, doc_id, filename, sha, texts): + chunks = [_chunk(t, i) for i, t in enumerate(texts)] + vectors = [embed(t) for t in texts] + store.create_document(conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id) + store.add_chunks(conn, scope, doc_id, chunks, vectors) + + +def test_lexical_returns_only_matching_docs(rag_conn): + _add_doc(rag_conn, "kb_a", "d1", "d1.txt", "h1", ["alpha bravo charlie"]) + _add_doc(rag_conn, "kb_a", "d2", "d2.txt", "h2", ["golf hotel india"]) + hits = store.search_lexical(rag_conn, "kb_a", "alpha", 10) + assert [cid for cid, _ in hits] == ["d1:0"] # d2 not returned (score 0) + + +def test_scope_isolation(rag_conn): + _add_doc(rag_conn, "kb_a", "d1", "f", "h1", ["alpha bravo"]) + _add_doc(rag_conn, "kb_b", "d2", "f", "h2", ["alpha bravo"]) + assert [cid for cid, _ in store.search_lexical(rag_conn, "kb_b", "alpha", 10)] == ["d2:0"] + + +def test_match_query_sanitizes_special_chars(): + assert store._match_query('AND OR "quote" (paren) -dash') != "" + + +def test_lexical_does_not_crash_on_punctuation(rag_conn): + _add_doc(rag_conn, "kb_a", "d1", "f", "h1", ["alpha bravo"]) + # Must not raise on FTS operators in the query. + store.search_lexical(rag_conn, "kb_a", 'NEAR("x" AND', 5) + + +def test_dense_ranks_by_cosine(rag_conn): + _add_doc(rag_conn, "kb_a", "d1", "f", "h1", ["alpha alpha"]) + _add_doc(rag_conn, "kb_a", "d2", "f", "h2", ["hotel golf"]) + ranked = store.search_dense(rag_conn, "kb_a", embed("alpha"), 10) + assert ranked[0][0] == "d1:0" and ranked[0][1] > 0.99 + + +def test_dense_empty_before_any_ingest(rag_conn): + # No chunks_vec table yet -> [], no crash. + assert store.search_dense(rag_conn, "kb_a", embed("alpha"), 10) == [] + + +def test_dedupe_by_hash(rag_conn): + _add_doc(rag_conn, "kb_a", "d1", "f", "SHA", ["alpha"]) + assert store.document_by_hash(rag_conn, "kb_a", "SHA") == "d1" + assert store.document_by_hash(rag_conn, "kb_a", "OTHER") is None + + +def test_delete_document_purges_all_tables(rag_conn): + _add_doc(rag_conn, "kb_a", "d1", "f", "h1", ["alpha bravo"]) + store.delete_document(rag_conn, "d1") + assert store.search_lexical(rag_conn, "kb_a", "alpha", 10) == [] + assert store.search_dense(rag_conn, "kb_a", embed("alpha"), 10) == [] + assert store.chunks_by_id(rag_conn, ["d1:0"]) == {} + assert store.get_document(rag_conn, "d1") is None + + +def test_incremental_add_is_flat(rag_conn): + # Adding doc2 must not touch doc1's fts rowids (append, not rebuild). + _add_doc(rag_conn, "kb_a", "d1", "f", "h1", ["alpha bravo charlie"]) + before = rag_conn.execute( + "SELECT rowid, chunk_id FROM chunks_fts WHERE scope='kb_a'" + ).fetchall() + _add_doc(rag_conn, "kb_a", "d2", "f", "h2", ["delta echo foxtrot"]) + after = rag_conn.execute( + "SELECT rowid, chunk_id FROM chunks_fts WHERE scope='kb_a' AND chunk_id LIKE 'd1:%'" + ).fetchall() + before_d1 = [(r["rowid"], r["chunk_id"]) for r in before if r["chunk_id"].startswith("d1:")] + after_d1 = [(r["rowid"], r["chunk_id"]) for r in after] + assert before_d1 == after_d1 + + +def test_chunks_by_id_joins_filename(rag_conn): + _add_doc(rag_conn, "kb_a", "d1", "paper.pdf", "h1", ["body text here"]) + rows = store.chunks_by_id(rag_conn, ["d1:0"]) + assert rows["d1:0"]["filename"] == "paper.pdf" + assert rows["d1:0"]["text"] == "body text here" + + +def test_kb_crud_and_delete_cascades(rag_conn): + kb_id = store.create_kb(rag_conn, name = "My KB", description = "d", kb_id = "K1") + assert store.get_kb(rag_conn, kb_id)["name"] == "My KB" + assert [k["id"] for k in store.list_kbs(rag_conn)] == ["K1"] + + scope = store.kb_scope("K1") + _add_doc(rag_conn, scope, "doc1", "f", "h1", ["alpha bravo"]) + store.delete_kb(rag_conn, "K1") + assert store.get_kb(rag_conn, "K1") is None + assert store.list_documents(rag_conn, scope) == [] + assert store.search_lexical(rag_conn, scope, "alpha", 10) == [] diff --git a/studio/backend/tests/test_recommended_folders_permission.py b/studio/backend/tests/test_recommended_folders_permission.py index 659c3b547d..33a457755e 100644 --- a/studio/backend/tests/test_recommended_folders_permission.py +++ b/studio/backend/tests/test_recommended_folders_permission.py @@ -6,17 +6,16 @@ Regression test for the /recommended-folders (and /browse-folders) 500 caused by an unreadable model directory, e.g. a stock root-owned ``ollama`` install at ``/usr/share/ollama/.ollama/models``. -Root cause: the folder-scan helpers in ``routes.models`` probed candidate -paths with a bare ``Path(p).is_dir()``. On Python <= 3.11 that returned -``False`` for an unreadable path; on Python >= 3.12 ``is_dir()`` propagates +Root cause: ``routes.models`` folder-scan helpers probed candidates with a +bare ``Path(p).is_dir()``. On Python <= 3.11 that returned ``False`` for an +unreadable path; on Python >= 3.12 ``is_dir()`` propagates ``PermissionError`` (EACCES), so the endpoint 500-ed through the whole -middleware stack instead of just skipping the directory. The probes now go -through the module-level ``_safe_is_dir`` helper. +middleware stack instead of skipping the directory. Probes now go through +the module-level ``_safe_is_dir`` helper. -``routes.models`` pulls the full backend dependency tree (fastapi, -structlog, the models package, ...), so rather than stand up the app we -extract the real ``_safe_is_dir`` definition from the source file and -exercise that exact function in isolation. The test therefore stays +``routes.models`` pulls the full backend dep tree (fastapi, structlog, the +models package, ...), so rather than stand up the app we extract the real +``_safe_is_dir`` from the source file and exercise it in isolation — dependency-free while still running the shipped code. Run: @@ -36,7 +35,7 @@ _models_src = _backend_root / "routes" / "models.py" def _load_safe_is_dir(): """Return the real ``_safe_is_dir`` from routes/models.py without - importing the (heavily dependency-laden) module.""" + importing the dependency-laden module.""" tree = ast.parse(_models_src.read_text()) fn = next( node @@ -51,8 +50,8 @@ def _load_safe_is_dir(): safe_is_dir = _load_safe_is_dir() -# Permission bits are bypassed for the superuser, so the chmod-000 setup -# below would not actually deny access when running as root. +# The superuser bypasses permission bits, so the chmod-000 setup below +# would not deny access when running as root. _skip_as_root = pytest.mark.skipif( hasattr(os, "geteuid") and os.geteuid() == 0, reason = "root bypasses filesystem permission bits", @@ -60,8 +59,7 @@ _skip_as_root = pytest.mark.skipif( def test_helper_exists_in_source(): - # Guards against a refactor silently dropping the helper the fix - # depends on (the extractor would then raise StopIteration). + # Guard against a refactor silently dropping the helper the fix needs. assert callable(safe_is_dir) @@ -83,8 +81,8 @@ def test_file_is_false(tmp_path): def test_mode000_dir_itself_is_still_a_dir(tmp_path): """A mode-000 directory is still stat-able via its (traversable) parent, so _safe_is_dir reports True without raising. Filtering out - dirs we cannot actually *read* is the caller's separate - os.access(R_OK|X_OK) check, not this helper's job.""" + dirs we can't *read* is the caller's separate os.access(R_OK|X_OK) + check, not this helper's job.""" locked = tmp_path / "locked" locked.mkdir() os.chmod(locked, 0o000) @@ -96,8 +94,8 @@ def test_mode000_dir_itself_is_still_a_dir(tmp_path): @_skip_as_root def test_path_under_unreadable_parent_returns_false_not_raises(tmp_path): - """The exact production scenario: stat()-ing a child of a mode-700 - system directory, e.g. ``/usr/share/ollama/.ollama/models``.""" + """The production scenario: stat()-ing a child of a mode-700 system + directory, e.g. ``/usr/share/ollama/.ollama/models``.""" parent = tmp_path / "ollama" parent.mkdir() os.chmod(parent, 0o000) @@ -113,8 +111,8 @@ def test_path_under_unreadable_parent_returns_false_not_raises(tmp_path): reason = "is_dir() only propagates PermissionError on Python >= 3.12", ) def test_demonstrates_the_underlying_stdlib_regression(tmp_path): - """Documents *why* _safe_is_dir exists: the old bare pattern raises - on the interpreters Studio ships on (3.12+).""" + """Documents *why* _safe_is_dir exists: the old bare pattern raises on + the interpreters Studio ships on (3.12+).""" parent = tmp_path / "ollama" parent.mkdir() os.chmod(parent, 0o000) diff --git a/studio/backend/tests/test_responses_api.py b/studio/backend/tests/test_responses_api.py index 5b55f87259..693e832113 100644 --- a/studio/backend/tests/test_responses_api.py +++ b/studio/backend/tests/test_responses_api.py @@ -1,18 +1,15 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -""" -Tests for the OpenAI Responses API schemas and input normalisation. -These tests do NOT require a running server or GPU -- they validate -the Pydantic models and the _normalise_responses_input helper. -""" +"""Tests for OpenAI Responses API Pydantic schemas and the +_normalise_responses_input helper. No server or GPU required.""" import sys import os import json import re -# Ensure backend is on path +# Ensure backend is on path. _backend = os.path.join(os.path.dirname(__file__), "..") sys.path.insert(0, _backend) @@ -33,34 +30,32 @@ from models.inference import ( ) -# ── _normalise_responses_input: copied from routes/inference.py ── -# We cannot import routes.inference directly because routes/__init__.py -# pulls in heavy dependencies (structlog/twisted/torch). This is a -# direct copy of the function for testing purposes. +# Copied from routes/inference.py: can't import it directly because +# routes/__init__.py pulls in heavy deps (structlog/twisted/torch). def _normalise_responses_input(payload: ResponsesRequest) -> list: - """Convert a ResponsesRequest into a list of ChatMessage for the completions backend.""" + """Convert a ResponsesRequest into ChatMessages for the completions backend.""" messages = [] - # System / developer instructions + # System / developer instructions. if payload.instructions: messages.append(ChatMessage(role = "system", content = payload.instructions)) - # Simple string input + # Simple string input. if isinstance(payload.input, str): if payload.input: messages.append(ChatMessage(role = "user", content = payload.input)) return messages - # List of ResponsesInputMessage + # List of ResponsesInputMessage. for msg in payload.input: role = "system" if msg.role == "developer" else msg.role if isinstance(msg.content, str): messages.append(ChatMessage(role = role, content = msg.content)) else: - # Convert Responses content parts -> Chat content parts + # Convert Responses content parts -> Chat content parts. parts = [] for part in msg.content: if isinstance(part, ResponsesInputTextPart): @@ -130,7 +125,7 @@ class TestResponsesRequest: assert req.instructions == "You are a helpful assistant." def test_extra_fields_accepted(self): - """OpenAI SDK may send fields we don't model -- extra='allow' should pass.""" + """OpenAI SDK may send unmodeled fields -- extra='allow' must pass.""" req = ResponsesRequest( input = "test", tools = [{"type": "web_search_preview"}], @@ -167,15 +162,13 @@ class TestResponsesRequest: class TestResponsesResponse: - """Validate response models serialise correctly.""" + """Response models serialise correctly.""" def test_basic_response(self): resp = ResponsesResponse( model = "test-model", output = [ - ResponsesOutputMessage( - content = [ResponsesOutputTextContent(text = "Hello!")] - ), + ResponsesOutputMessage(content = [ResponsesOutputTextContent(text = "Hello!")]), ], usage = ResponsesUsage(input_tokens = 10, output_tokens = 5, total_tokens = 15), ) @@ -226,7 +219,7 @@ class TestResponsesResponse: class TestNormaliseResponsesInput: - """Test _normalise_responses_input converts Responses input to ChatMessages.""" + """_normalise_responses_input converts Responses input to ChatMessages.""" def test_string_input(self): payload = ResponsesRequest(input = "Hello world") @@ -324,5 +317,4 @@ class TestNormaliseResponsesInput: if __name__ == "__main__": import pytest - pytest.main([__file__, "-v"]) diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 2f1161c329..69fb0a78c2 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -6,19 +6,19 @@ Tests for the OpenAI /v1/responses client-side function-calling pass-through. Covers: - ResponsesRequest accepts Responses-shape `tools`, `tool_choice`, - `parallel_tool_calls`, and the `function_call` / `function_call_output` - input items used for multi-turn tool loops. -- _translate_responses_tools_to_chat() converts the flat Responses tool - shape to the nested Chat Completions shape, drops non-function built-in - tools, and returns None for empty lists. -- _translate_responses_tool_choice_to_chat() passes string choices through - and converts {type:function,name:X} to Chat Completions' nested shape. -- _normalise_responses_input() maps function_call_output items to + `parallel_tool_calls`, and `function_call` / `function_call_output` + input items for multi-turn tool loops. +- _translate_responses_tools_to_chat(): flat Responses tool shape -> + nested Chat Completions shape, drops non-function built-in tools, + returns None for empty lists. +- _translate_responses_tool_choice_to_chat(): passes string choices + through, converts {type:function,name:X} to the nested shape. +- _normalise_responses_input(): maps function_call_output items to role="tool" ChatMessages with tool_call_id, and function_call items to assistant messages with tool_calls. -- _chat_tool_calls_to_responses_output() preserves call_id and drops +- _chat_tool_calls_to_responses_output(): keeps call_id, drops non-function tool calls. -- ResponsesOutputFunctionCall and ResponsesResponse round-trip tool-call +- ResponsesOutputFunctionCall / ResponsesResponse round-trip tool-call outputs without losing fields. No running server or GPU required. @@ -26,12 +26,15 @@ No running server or GPU required. import os import sys +import asyncio +from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") sys.path.insert(0, _backend) import json +import httpx import pytest from pydantic import ValidationError @@ -52,8 +55,11 @@ from models.inference import ( ResponsesUsage, ) from routes.inference import ( + _build_chat_request, _chat_tool_calls_to_responses_output, _normalise_responses_input, + _responses_tool_output_text, + _responses_stream, _translate_responses_tool_choice_to_chat, _translate_responses_tools_to_chat, ) @@ -104,9 +110,9 @@ class TestResponsesRequestTools: assert req.parallel_tool_calls is True def test_builtin_tool_type_passes_validation(self): - """Non-function built-in tools (web_search, file_search, mcp, ...) must - not raise at request validation so SDKs that default to them don't - fail on Studio; they are filtered out during translation.""" + """Non-function built-in tools (web_search, file_search, mcp, ...) + must not raise at validation so SDKs that default to them don't + fail on Studio; they're filtered out during translation.""" req = ResponsesRequest( input = "hi", tools = [{"type": "web_search_preview"}], @@ -160,9 +166,7 @@ class TestResponsesMultiTurnInput: def test_function_call_output_missing_call_id_rejected(self): with pytest.raises(ValidationError): - ResponsesFunctionCallOutputInputItem( - type = "function_call_output", output = "x" - ) + ResponsesFunctionCallOutputInputItem(type = "function_call_output", output = "x") def test_function_call_output_accepts_content_array(self): item = ResponsesFunctionCallOutputInputItem( @@ -222,9 +226,7 @@ class TestToolsTranslation: assert _translate_responses_tools_to_chat([]) is None def test_only_builtin_tools_returns_none(self): - assert ( - _translate_responses_tools_to_chat([{"type": "web_search_preview"}]) is None - ) + assert _translate_responses_tools_to_chat([{"type": "web_search_preview"}]) is None def test_description_optional(self): out = _translate_responses_tools_to_chat( @@ -253,18 +255,36 @@ class TestToolChoiceTranslation: ) == {"type": "function", "function": {"name": "get_weather"}} def test_already_chat_nested_shape_passes_through(self): - """If a client happens to send the Chat Completions nested shape, - we don't double-wrap it.""" + """A client sending the Chat Completions nested shape isn't + double-wrapped.""" already_nested = {"type": "function", "function": {"name": "get_weather"}} - assert ( - _translate_responses_tool_choice_to_chat(already_nested) == already_nested - ) + assert _translate_responses_tool_choice_to_chat(already_nested) == already_nested def test_unknown_shape_passes_through(self): obj = {"type": "allowed_tools", "tools": [{"type": "function", "name": "x"}]} assert _translate_responses_tool_choice_to_chat(obj) == obj +class TestBuildChatRequest: + def test_parallel_tool_calls_false_is_preserved_for_passthrough_caps(self): + payload = ResponsesRequest( + input = "hi", + tools = [ + { + "type": "function", + "name": "lookup", + "parameters": {"type": "object"}, + } + ], + parallel_tool_calls = False, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = True) + + assert chat_req.parallel_tool_calls is False + + # ===================================================================== # _normalise_responses_input — multi-turn tool mapping # ===================================================================== @@ -303,10 +323,10 @@ class TestNormaliseResponsesInputWithTools: def test_instructions_plus_developer_message_are_merged(self): """Codex CLI sends `instructions` (system prompt) AND a developer - message in `input`. Strict chat templates (harmony / gpt-oss, Qwen3, - ...) raise "System message must be at the beginning" when two - separate system-role messages appear, so we must emit exactly one - merged system message at the top. + message in `input`. Strict chat templates (harmony / gpt-oss, + Qwen3, ...) raise "System message must be at the beginning" on two + separate system-role messages, so we emit exactly one merged + system message at the top. """ payload = ResponsesRequest( instructions = "Base instructions.", @@ -320,14 +340,14 @@ class TestNormaliseResponsesInputWithTools: assert len(system_roles) == 1 assert "Base instructions." in system_roles[0].content assert "Developer override." in system_roles[0].content - # System must be the very first message for strict templates. + # System must be the first message for strict templates. assert msgs[0].role == "system" assert msgs[1].role == "user" def test_developer_message_after_user_is_still_hoisted(self): - """Multi-turn conversations where a developer message appears after - user turns must still produce a single leading system message, not - a mid-conversation system that strict templates reject.""" + """A developer message appearing after user turns must still + produce a single leading system message, not a mid-conversation + system that strict templates reject.""" payload = ResponsesRequest( input = [ {"role": "user", "content": "Hello"}, @@ -374,6 +394,100 @@ class TestNormaliseResponsesInputWithTools: # Content is serialised so llama-server sees a string. assert json.loads(msgs[0].content) == [{"type": "output_text", "text": "ok"}] + def test_empty_function_call_output_gets_no_output_sentinel(self): + payload = ResponsesRequest( + input = [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": "", + } + ], + ) + msgs = _normalise_responses_input(payload) + assert msgs[0].role == "tool" + assert msgs[0].tool_call_id == "call_1" + assert msgs[0].content == "(no output)" + ChatMessage(**msgs[0].model_dump(exclude_none = True)) + + def test_whitespace_function_call_output_gets_no_output_sentinel(self): + payload = ResponsesRequest( + input = [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": " \n\t", + } + ], + ) + msgs = _normalise_responses_input(payload) + assert msgs[0].content == "(no output)" + + def test_empty_content_array_output_gets_no_output_sentinel(self): + payload = ResponsesRequest( + input = [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": [], + } + ], + ) + msgs = _normalise_responses_input(payload) + assert msgs[0].content == "(no output)" + + def test_image_content_array_tool_output_is_serialised(self): + payload = ResponsesRequest( + input = [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "iVBORw0KGgo=", + }, + } + ], + } + ], + ) + msgs = _normalise_responses_input(payload) + assert msgs[0].role == "tool" + assert json.loads(msgs[0].content)[0]["type"] == "image" + + def test_image_payload_outside_output_gets_no_output_sentinel(self): + payload = ResponsesRequest( + input = [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": "", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "iVBORw0KGgo=", + }, + } + ], + } + ], + ) + msgs = _normalise_responses_input(payload) + assert msgs[0].role == "tool" + assert msgs[0].tool_call_id == "call_1" + assert msgs[0].content == "(no output)" + + def test_tool_output_serializer_preserves_non_empty_text(self): + assert _responses_tool_output_text("done") == "done" + assert _responses_tool_output_text(" done ") == " done " + # ===================================================================== # Response mapping — tool_calls → function_call output items @@ -430,6 +544,134 @@ class TestChatToolCallsToResponsesOutput: assert items[0]["arguments"] == "" +# ===================================================================== +# Streaming Responses adapter +# ===================================================================== + + +class TestResponsesStreamAdapter: + class _Request: + async def is_disconnected(self): + return False + + @staticmethod + async def _collect(response): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk) + return chunks + + @staticmethod + def _payloads(lines, event_name): + prefix = f"event: {event_name}\n" + return [ + json.loads(line.split("data: ", 1)[1].strip()) + for line in lines + if line.startswith(prefix) + ] + + def test_requests_usage_and_caps_parallel_tool_calls(self, monkeypatch): + import routes.inference as inf_mod + + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode()) + chunks = [ + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "type": "function", + "function": {"name": "first", "arguments": "{}"}, + }, + { + "index": 1, + "id": "call_1", + "type": "function", + "function": {"name": "second", "arguments": "{}"}, + }, + ] + } + } + ] + }, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + content = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + content += "data: [DONE]\n\n" + return httpx.Response( + 200, + content = content.encode(), + headers = {"content-type": "text/event-stream"}, + ) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def _client(*args, **kwargs): + return real_async_client( + transport = transport, + timeout = kwargs.get("timeout", 600), + ) + + monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + is_vision = False, + context_length = 4096, + base_url = "http://llama.test", + # Non-reasoning template: the real backend returns None here. + _request_reasoning_kwargs = ( + lambda enable_thinking = None, reasoning_effort = None, preserve_thinking = None: None + ), + ), + ) + + payload = ResponsesRequest( + input = "hi", + stream = True, + parallel_tool_calls = False, + tools = [ + { + "type": "function", + "name": "first", + "parameters": {"type": "object"}, + }, + { + "type": "function", + "name": "second", + "parameters": {"type": "object"}, + }, + ], + ) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + assert captured["body"]["stream_options"] == {"include_usage": True} + joined = "".join(lines) + assert "call_0" in joined + assert "call_1" not in joined + completed = self._payloads(lines, "response.completed")[0] + assert completed["response"]["usage"] == { + "input_tokens": 2, + "output_tokens": 3, + "total_tokens": 5, + } + + # ===================================================================== # Response model — ResponsesOutputFunctionCall / mixed output # ===================================================================== @@ -492,8 +734,8 @@ class TestCodexStyleRequestShapes: """Regression tests for the request shapes OpenAI Codex CLI sends.""" def test_assistant_replay_output_text_accepted(self): - """Codex replays prior assistant turns with `output_text` content. - Before, this triggered a 422 on every turn after the first.""" + """Codex replays prior assistant turns with `output_text` content; + this used to 422 on every turn after the first.""" req = ResponsesRequest( input = [ {"role": "user", "content": "Hi"}, @@ -520,7 +762,7 @@ class TestCodexStyleRequestShapes: def test_reasoning_item_accepted_as_unknown(self): """`reasoning` items replayed from prior o-series turns must not - fail validation — Codex preserves them in multi-turn.""" + fail validation — Codex keeps them in multi-turn.""" req = ResponsesRequest( input = [ {"role": "user", "content": "Hi"}, @@ -538,7 +780,7 @@ class TestCodexStyleRequestShapes: def test_unknown_content_part_type_accepted(self): """Unknown content-part types (e.g. future input_audio) validate as - ResponsesUnknownContentPart so the whole request doesn't 422.""" + ResponsesUnknownContentPart so the request doesn't 422.""" req = ResponsesRequest( input = [ { @@ -602,15 +844,15 @@ class TestCodexStyleRequestShapes: ], ) msgs = _normalise_responses_input(payload) - # Single leading merged system; no mid-conversation system. + # One leading merged system; no mid-conversation system. assert msgs[0].role == "system" assert sum(1 for m in msgs if m.role == "system") == 1 assert "Base instructions." in msgs[0].content assert "Dev override." in msgs[0].content roles = [m.role for m in msgs[1:]] - # Reasoning item is dropped. Order: user, assistant(tool_calls), - # tool, assistant(text), user. + # Reasoning dropped. Order: user, assistant(tool_calls), tool, + # assistant(text), user. assert roles == ["user", "assistant", "tool", "assistant", "user"] assert msgs[2].tool_calls is not None assert msgs[3].role == "tool" @@ -618,16 +860,14 @@ class TestCodexStyleRequestShapes: assert msgs[4].content == "It's 20°C." def test_single_output_text_part_flattens_to_string(self): - """ChatMessage assistant role prefers plain string content — tests - confirm we don't forward a single-part array that would otherwise - force legacy chat templates into multimodal handling.""" + """ChatMessage assistant role prefers plain string content — we + don't forward a single-part array that would force legacy chat + templates into multimodal handling.""" payload = ResponsesRequest( input = [ { "role": "assistant", - "content": [ - {"type": "output_text", "text": "ok", "annotations": []} - ], + "content": [{"type": "output_text", "text": "ok", "annotations": []}], }, {"role": "user", "content": "next"}, ], @@ -638,9 +878,9 @@ class TestCodexStyleRequestShapes: class TestTranslatedMessagesValidate: - """Verify that the messages produced by _normalise_responses_input - satisfy ChatMessage's role-shape validator so the downstream /v1/chat/ - completions pass-through does not reject them.""" + """Messages from _normalise_responses_input satisfy ChatMessage's + role-shape validator so the downstream /v1/chat/completions + pass-through doesn't reject them.""" def test_round_trip_multi_turn(self): payload = ResponsesRequest( @@ -662,6 +902,20 @@ class TestTranslatedMessagesValidate: ) msgs = _normalise_responses_input(payload) for m in msgs: - # Constructing a fresh ChatMessage from the dump round-trips the - # role-shape validator — the key invariant for the passthrough. + # Building a fresh ChatMessage from the dump round-trips the + # role-shape validator — the passthrough's key invariant. + ChatMessage(**m.model_dump(exclude_none = True)) + + def test_empty_tool_output_round_trips_through_chat_message_validator(self): + payload = ResponsesRequest( + input = [ + { + "type": "function_call_output", + "call_id": "call_empty", + "output": "", + }, + ], + ) + msgs = _normalise_responses_input(payload) + for m in msgs: ChatMessage(**m.model_dump(exclude_none = True)) diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py new file mode 100644 index 0000000000..6e70c7cde4 --- /dev/null +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -0,0 +1,243 @@ +# 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 _rocm_classify_unified_memory (ROCm OOM-guard classifier). + +Three paths: (1) canonical gcnArchName, (2) alternate-spelling attr, (3) all +arch attrs absent -> device-name substring match. + +Regression: Strix Halo (gfx1151) was misclassified as discrete on Radeon wheels +that set props.name="Radeon 8060S Graphics" but no gcnArchName, applying the +wrong headroom factor on a 128 GiB unified-memory pool. +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from core.training.worker import _rocm_classify_unified_memory + + +# ── helpers ────────────────────────────────────────────────────────────────── + + +def _props(**kwargs) -> SimpleNamespace: + """Build a fake device-properties object with the given attributes.""" + return SimpleNamespace(**kwargs) + + +# ── Path 0: props.is_integrated (driver's own unified-memory answer) ───────── + + +class TestIsIntegratedSignal: + """hipDeviceProp_t.integrated wins when truthy; 0/absent never downgrades. + + Same universal gate PR #5988's UMA safetensors fast-load uses -- keeps + Studio's two unified-memory consumers on one signal.""" + + def test_integrated_upgrades_unknown_apu(self) -> None: + # gfx1103 Phoenix iGPU: outside the hardcoded arch set, but the + # driver says integrated -> unified. + props = _props(gcnArchName = "gfx1103", name = "Radeon 780M", is_integrated = 1) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1103" + assert is_unified is True + + def test_integrated_wins_without_any_arch(self) -> None: + props = _props(name = "Some Future APU", is_integrated = 1) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert is_unified is True + + def test_zero_does_not_downgrade_known_apu(self) -> None: + # A wheel that zeroes the field must not flip Strix Halo to discrete. + props = _props(gcnArchName = "gfx1151", name = "x", is_integrated = 0) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert is_unified is True + + def test_absent_keeps_existing_behavior(self) -> None: + props = _props(gcnArchName = "gfx1201", name = "RX 9070 XT") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert is_unified is False + + def test_discrete_with_zero_stays_discrete(self) -> None: + props = _props(gcnArchName = "gfx1100", name = "RX 7900 XTX", is_integrated = 0) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert is_unified is False + + +# ── Path 1: canonical gcnArchName ──────────────────────────────────────────── + + +class TestCanonicalGcnArchName: + """gcnArchName is present and populated.""" + + @pytest.mark.parametrize( + "arch, expected_unified", + [ + ("gfx1150", True), # Strix Point + ("gfx1151", True), # Strix Halo + ("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete + ("gfx906", False), # MI50 — discrete server GPU + ("gfx1201", False), # RX 9070 XT — discrete + ], + ) + def test_canonical_attr(self, arch: str, expected_unified: bool) -> None: + props = _props(gcnArchName = arch, name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == arch + assert is_unified is expected_unified + + def test_arch_with_colon_suffix_stripped(self) -> None: + """gcnArchName can carry xnack/sramecc suffix; only the base is kept.""" + props = _props(gcnArchName = "gfx1151:xnack-", name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + def test_canonical_attr_wins_over_name(self) -> None: + """Arch attr takes priority; device name is ignored.""" + # Discrete arch, but name looks like a unified SKU — arch must win. + props = _props(gcnArchName = "gfx1100", name = "Radeon 890M") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1100" + assert is_unified is False + + +# ── Path 2: alternate-spelling fallback ────────────────────────────────────── + + +class TestAlternateSpellingFallback: + """gcnArchName is missing but an alternate attr spelling is present.""" + + @pytest.mark.parametrize( + "attr_name", + ["gcn_arch_name", "arch_name", "gfx_arch_name"], + ) + def test_alternate_attr_unified(self, attr_name: str) -> None: + props = _props(**{attr_name: "gfx1151"}, name = "Radeon 8060S Graphics") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + @pytest.mark.parametrize( + "attr_name", + ["gcn_arch_name", "arch_name", "gfx_arch_name"], + ) + def test_alternate_attr_discrete(self, attr_name: str) -> None: + props = _props(**{attr_name: "gfx1201"}, name = "Radeon RX 9070 XT") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1201" + assert is_unified is False + + def test_first_non_empty_attr_wins(self) -> None: + """With multiple alternate attrs, the first non-empty one wins.""" + props = _props(gcn_arch_name = "gfx1151", arch_name = "gfx1100", name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + +# ── Path 3: device-name fallback ───────────────────────────────────────────── + + +class TestDeviceNameFallback: + """ALL arch attrs absent — classifier must rely solely on device name.""" + + # --- unified-memory devices that MUST be detected --- + + @pytest.mark.parametrize( + "device_name", + [ + # gfx1150 Strix Point + "Radeon 890M", + "AMD Radeon 890M Graphics", + "RADEON 890M", # case-insensitive + "Radeon 880M", + "AMD Radeon 880M Graphics", + # gfx1151 Strix Halo — the regression case from the review + "Radeon 8060S Graphics", # Ryzen AI MAX+ 395 (as returned by torch) + "AMD Radeon 8060S", + "Radeon 8050S Graphics", # cut-down Strix Halo SKU + "AMD Radeon 8050S", + # case variants + "RADEON 8060S GRAPHICS", + "radeon 8050s", + ], + ) + def test_unified_memory_detected(self, device_name: str) -> None: + props = _props(name = device_name) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "", f"expected empty gcn_arch, got {gcn!r}" + assert is_unified is True, f"device {device_name!r} should be classified as unified-memory" + + # --- discrete devices that must NOT be mis-classified --- + + @pytest.mark.parametrize( + "device_name", + [ + "Radeon RX 9070 XT", + "AMD Radeon RX 7900 XTX", + "Radeon RX 6900 XT", + "Radeon Pro W7900", + "AMD Instinct MI300X", + # Superficially similar substrings but discrete + "Radeon RX 580", + "Radeon VII", + ], + ) + def test_discrete_not_misclassified(self, device_name: str) -> None: + props = _props(name = device_name) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert ( + is_unified is False + ), f"discrete device {device_name!r} should NOT be classified as unified-memory" + + def test_empty_name_returns_false(self) -> None: + """Absent name must not crash and must default to discrete.""" + props = _props() # no 'name' attr at all + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert is_unified is False + + def test_none_name_returns_false(self) -> None: + props = _props(name = None) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert is_unified is False + + +# ── Fraction selection (source-pinned) ─────────────────────────────────────── + + +_WORKER_PY = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" + + +class TestMemFractionSelection: + """Pin the per-platform fraction policy in worker.py section 1g. + + On native Windows, torch.cuda.mem_get_info's total is the WDDM budget + the driver grants HIP -- the OS share of RAM is already outside it, so + a 0.80 cap double-taxes (field report: 48.49 GiB budget -> '38.79 GiB + allowed' OOM denying a 47.29 GiB load that fit in free memory). 1.0 + removes the double-tax; current AMD Windows wheels enforce only + sub-1.0 fractions, so it behaves like torch's uncapped default with + WDDM arbitrating residency (measured on gfx1151).""" + + def test_unified_win32_uses_budget_exact_fraction(self) -> None: + source = _WORKER_PY.read_text(encoding = "utf-8") + assert '1.0 if sys.platform == "win32" else 0.80' in source + + def test_discrete_keeps_090(self) -> None: + source = _WORKER_PY.read_text(encoding = "utf-8") + assert "_mem_fraction = 0.90" in source + + def test_win32_unified_logs_vgm_hint(self) -> None: + """Users must learn the WDDM budget is raisable (BIOS UMA / AMD + Software Variable Graphics Memory) instead of assuming a bug.""" + source = _WORKER_PY.read_text(encoding = "utf-8") + assert "Variable Graphics Memory" in source diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index c3ee5b9ff1..1e8fb9e2b2 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -1,11 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Capability advertisement contract: classifier honesty, worker→ -orchestrator IPC hop, and route-layer end-to-end. Pure helpers + fakes; -no torch / transformers import. -""" +"""Capability advertisement contract: classifier honesty, worker->orchestrator +IPC hop, route-layer end-to-end. Pure helpers + fakes; no torch/transformers.""" from __future__ import annotations @@ -116,7 +113,7 @@ def test_detect_safetensors_features_none_template_returns_all_false(): def test_detect_safetensors_features_gptoss_disables_tools(): - """gpt-oss Harmony: tools intentionally off even if template marks it.""" + """gpt-oss Harmony: tools off even if template marks it.""" from routes.inference import _detect_safetensors_features backend = MagicMock() @@ -129,11 +126,9 @@ def test_detect_safetensors_features_gptoss_disables_tools(): assert flags["supports_tools"] is False -# Llama-3 / Mistral templates advertise tool handling but the model emits -# tool calls in <|python_tag|> / [TOOL_CALLS] format -- not the -# / / [TOOL_CALLS], +# which our parser can't read. The route helper must not flip supports_tools=True +# for them, else the UI enables a pill the agentic loop can't honour. LLAMA3_TEMPLATE = """ {%- if tools %} @@ -207,9 +202,8 @@ def test_detect_safetensors_features_function_xml_format_keeps_tools_on(): assert flags["supports_tools"] is True -# Qwen3.5 family pins -- the live GGUF + safetensors templates fetched -# from the unsloth/Qwen3.5-0.8B(-GGUF) repos both wrap tool calls as -# ``\n...``. Capture a faithful slice so the +# Qwen3.5 family pin: the live GGUF + safetensors templates both wrap tool +# calls as ``\n...``. Faithful slice so the # classifier never silently regresses for this family. QWEN35_TOOL_INSTRUCTION = ( @@ -234,7 +228,7 @@ QWEN35_TOOL_INSTRUCTION = ( def test_detect_safetensors_features_qwen35_keeps_tools_on(): - """unsloth/Qwen3.5-0.8B family must surface tools+reasoning enabled.""" + """unsloth/Qwen3.5-0.8B family must surface tools+reasoning on.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Qwen3.5-0.8B") @@ -248,7 +242,7 @@ def test_detect_safetensors_features_qwen35_keeps_tools_on(): def test_orchestrator_mirrors_chat_template_info_into_models_dict(): - """Worker → orchestrator must copy chat_template_info verbatim.""" + """Worker → orchestrator copies chat_template_info verbatim.""" from core.inference.orchestrator import InferenceOrchestrator orch = InferenceOrchestrator.__new__(InferenceOrchestrator) @@ -274,7 +268,7 @@ def test_orchestrator_mirrors_chat_template_info_into_models_dict(): }, } - # Replay orchestrator.load_model's mirror block verbatim. + # Replay orchestrator.load_model's mirror block. orch.active_model_name = model_info["identifier"] orch.models[orch.active_model_name] = { "is_vision": model_info.get("is_vision", False), @@ -369,11 +363,7 @@ def test_worker_load_reply_payload_includes_chat_template_info(): "is_gguf": False, } _bm = getattr(backend, "models", {}) or {} - _entry = ( - _bm.get(mc.identifier) - or _bm.get(getattr(backend, "active_model_name", None)) - or {} - ) + _entry = _bm.get(mc.identifier) or _bm.get(getattr(backend, "active_model_name", None)) or {} _tpl_info = _entry.get("chat_template_info") if isinstance(_tpl_info, dict): model_info["chat_template_info"] = { @@ -390,7 +380,7 @@ def test_worker_load_reply_payload_includes_chat_template_info(): def test_worker_load_reply_payload_survives_missing_template(): - """Tokenizer with no chat_template still produces a valid reply.""" + """Tokenizer with no chat_template still yields a valid reply.""" class _StubBackend: def __init__(self): @@ -425,7 +415,7 @@ def test_worker_load_reply_payload_survives_missing_template(): def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors(): - """End-to-end: Qwen3 safetensors flips supports_tools=True.""" + """E2E: Qwen3 safetensors flips supports_tools=True.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace( diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 923af87c4f..8aa6e5df4e 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -1,42 +1,29 @@ # 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 the safetensors agentic tool loop. +"""Tests for the safetensors agentic tool loop. -Covers the shared ``tool_call_parser`` helpers and the cumulative-text -state machine inside ``safetensors_agentic.run_safetensors_tool_loop``. -The loop is exercised with hand-crafted fake single-turn generators so -no model load is needed; the tests run in CI under a few seconds. - -Edge cases under coverage: -* Plain answers (no tool calls) flush full content. -* Single ``{json}`` triggers the tool and re-enters. -* Single ``...`` XML form triggers the same path. -* Truncated unclosed ```` is still parsed. -* Tool result is fed back as ``role=tool`` for the next iteration. -* Bad JSON inside ```` does not raise and (when healed) is - routed as a ``{"query": ...}`` web search call. -* Duplicate tool calls produce a synthetic "do not repeat" result the - second time. -* ``__IMAGES__`` sentinel is stripped before the model sees the result. -* Tool execution errors are tagged so the model gets a nudge but the - loop keeps streaming. -* Cancel is honoured between iterations. -* ``max_tool_iterations`` cap is respected and a final-answer attempt - closes the stream cleanly. +Covers the ``tool_call_parser`` helpers and the cumulative-text state machine in +``run_safetensors_tool_loop``, run against fake single-turn generators (no model +load). Edge cases: plain answers, JSON and XML tool-call forms, truncated/unclosed +calls, tool-result feedback, bad-JSON heal, duplicate-call short-circuit, +``__IMAGES__`` sentinel stripping, executor errors, cancel, and the iteration cap. """ import threading +from typing import cast import pytest from core.inference import safetensors_agentic from core.inference.safetensors_agentic import ( _coerce_arguments, + _detect_render_html_tool_start, run_safetensors_tool_loop, + strip_tool_markup_streaming, ) from core.inference.tool_call_parser import ( + RAG_MAX_SEARCHES_PER_TURN, has_tool_signal, parse_tool_calls_from_text, strip_tool_markup, @@ -51,9 +38,7 @@ from utils.datasets import is_gpt_oss_model_name class TestParser: def test_json_tool_call(self): - text = ( - '{"name":"web_search","arguments":{"query":"hello"}}' - ) + text = '{"name":"web_search","arguments":{"query":"hello"}}' result = parse_tool_calls_from_text(text) assert len(result) == 1 tc = result[0] @@ -64,12 +49,17 @@ class TestParser: assert "hello" in tc["function"]["arguments"] def test_json_tool_call_unclosed(self): - # No ; balanced-brace extractor must still close. + # No ; balanced-brace extractor must still close it. text = '{"name":"python","arguments":{"code":"print(1)"}}' result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "python" + def test_json_tool_call_unclosed_requires_healing(self): + text = '{"name":"python","arguments":{"code":"print(1)"}}' + assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + def test_xml_function_call(self): text = "print('hi')" result = parse_tool_calls_from_text(text) @@ -85,10 +75,14 @@ class TestParser: assert result[0]["function"]["name"] == "terminal" assert "ls -la" in result[0]["function"]["arguments"] + def test_xml_unclosed_requires_healing(self): + text = "ls -la" + assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "terminal" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + def test_code_with_embedded_xml(self): - # A code parameter contains the literal . Must not - # truncate the value because the parser uses end-of-body as the - # only boundary for single-parameter calls. + # A code parameter with a literal must not truncate: the + # parser uses end-of-body as the only boundary for single-param calls. text = ( "html = ''\n" "print('hi')" @@ -97,6 +91,17 @@ class TestParser: assert len(result) == 1 assert "print('hi')" in result[0]["function"]["arguments"] + def test_function_signal_inside_parameter_is_literal(self): + text = ( + "" + "print('')" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "" in result[0]["function"]["arguments"] + def test_multiple_calls(self): text = ( '{"name":"web_search","arguments":{"query":"a"}}' @@ -110,7 +115,7 @@ class TestParser: def test_bad_json_does_not_raise(self): text = "{not valid json}" result = parse_tool_calls_from_text(text) - # Bad JSON is silently dropped; caller can fall back to text. + # Bad JSON is dropped silently; caller can fall back to text. assert result == [] def test_has_tool_signal(self): @@ -118,17 +123,46 @@ class TestParser: assert has_tool_signal("hi ...") assert not has_tool_signal("hello world") + def test_render_html_start_detector_uses_first_tool(self): + assert _detect_render_html_tool_start("") + assert _detect_render_html_tool_start( + '{"name":"render_html","arguments":{"code":""}' + ) + assert not _detect_render_html_tool_start( + "''" + ) + assert not _detect_render_html_tool_start( + '{"name":"python","arguments":{"code":""}}' + ) + def test_strip_markup_closed(self): text = "before {} after" assert strip_tool_markup(text) == "before after" def test_strip_markup_unclosed_final(self): text = "before {partial" - # With final=True the trailing run is dropped. + # final=True drops the trailing run. assert strip_tool_markup(text, final = True) == "before" # Without final=True the unclosed run is preserved. assert "partial" in strip_tool_markup(text) + def test_streaming_strip_respects_disabled_healing(self): + raw = 'before {"name":"web_search"' + assert strip_tool_markup_streaming(raw, auto_heal_tool_calls = False) == raw + assert strip_tool_markup_streaming(raw) == "before " + + def test_streaming_strip_respects_disabled_healing_without_tool_protocol(self): + raw = 'before {"name":"web_search"' + assert strip_tool_markup_streaming(raw, auto_heal_tool_calls = False) == raw + assert ( + strip_tool_markup_streaming( + raw, + auto_heal_tool_calls = False, + tool_protocol_active = True, + ) + == "before " + ) + # ──────────────────────────────────────────────────────────────────── # run_safetensors_tool_loop @@ -172,6 +206,7 @@ class FakeExecuteTool: cancel_event = None, timeout = None, session_id = None, + rag_scope = None, ): self.calls.append((name, arguments)) result = self.results.pop(0) if self.results else "OK" @@ -189,11 +224,15 @@ def _collect_events(generator, max_events = 200): return events -def _make_loop(*, turns, exec_results = None, **kwargs): +def _make_loop( + *, + turns, + exec_results = None, + **kwargs, +): """Build a configured loop with a multi-turn fake generator. - ``turns`` is a list of chunk-lists; iteration N yields chunks from - ``turns[N]``. + ``turns`` is a list of chunk-lists; iteration N yields chunks from ``turns[N]``. """ turn_iter = iter(turns) @@ -221,6 +260,41 @@ def _make_loop(*, turns, exec_results = None, **kwargs): ), exec_fn +def test_active_tools_are_passed_to_single_turn_after_render_html_success(): + captured_tool_names: list[list[str]] = [] + exec_fn = FakeExecuteTool(["Rendered HTML artifact."]) + + def fake_single_turn(_messages, *, active_tools = None): + captured_tool_names.append( + [ + (tool.get("function") or {}).get("name") + for tool in (active_tools or []) + if (tool.get("function") or {}).get("name") + ] + ) + if len(captured_tool_names) == 1: + yield '{"name":"render_html","arguments":{"code":"one"}}' + else: + yield "Done." + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "make html"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "web_search"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + ) + + assert exec_fn.calls == [("render_html", {"code": "one"})] + assert captured_tool_names == [["render_html", "web_search"], ["web_search"]] + assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) + + class TestLoopBasic: def test_plain_answer(self): # No tool XML; loop should yield content then status="". @@ -232,7 +306,7 @@ class TestLoopBasic: contents = [e for e in events if e["type"] == "content"] statuses = [e for e in events if e["type"] == "status"] assert contents, "expected at least one content event" - # Final cumulative content should contain the answer. + # Final cumulative content must contain the answer. final_text = contents[-1]["text"] assert "Hello world!" in final_text assert statuses and statuses[-1]["text"] == "" @@ -240,13 +314,13 @@ class TestLoopBasic: def test_single_tool_then_answer(self): loop, exec_fn = _make_loop( turns = [ - # : tool call only. + # Tool call only. [ '{"name":"web_search",', '"arguments":{"query":"weather"}}', "", ], - # : final answer. + # Final answer. ["The ", "weather is ", "sunny."], ], exec_results = ["Sunny and 22C"], @@ -256,7 +330,7 @@ class TestLoopBasic: assert "tool_start" in kinds assert "tool_end" in kinds - # Tool was actually called with the parsed arguments. + # Tool was called with the parsed arguments. assert exec_fn.calls == [("web_search", {"query": "weather"})] tool_start = next(e for e in events if e["type"] == "tool_start") @@ -280,11 +354,102 @@ class TestLoopBasic: contents = [e for e in events if e["type"] == "content"] assert "Result: 1" in contents[-1]["text"] + def test_render_html_emits_provisional_tool_start(self): + exec_fn = FakeExecuteTool(["Rendered HTML artifact."]) + turn_iter = iter( + [ + [ + "", + "", + "Hi", + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "make html"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + execute_tool = exec_fn, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + assert len(tool_starts) == 2 + assert tool_starts[0]["tool_name"] == "render_html" + assert tool_starts[0]["arguments"] == {} + assert tool_starts[1]["tool_name"] == "render_html" + assert "" in tool_starts[1]["arguments"]["code"] + assert exec_fn.calls[0][0] == "render_html" + assert "" in exec_fn.calls[0][1]["code"] + + def test_python_tool_containing_render_html_signal_does_not_emit_provisional_start(self): + loop, exec_fn = _make_loop( + turns = [ + [ + "", + "print('')", + "", + ], + ["Done."], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + assert len(tool_starts) == 1 + assert tool_starts[0]["tool_name"] == "python" + assert exec_fn.calls == [("python", {"code": "print('')"})] + + def test_render_html_success_blocks_second_artifact_call(self): + exec_fn = FakeExecuteTool(["Rendered HTML artifact."]) + turn_iter = iter( + [ + [ + '{"name":"render_html",', + '"arguments":{"code":"one"}}', + ], + [ + '{"name":"render_html",', + '"arguments":{"code":"two"}}', + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "make html"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + execute_tool = exec_fn, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + assert exec_fn.calls == [("render_html", {"code": "one"})] + assert [e["arguments"] for e in tool_starts] == [{}, {"code": "one"}] + def test_truncated_unclosed_tool_call(self): loop, exec_fn = _make_loop( turns = [ - # No ; balanced-brace parser must still - # succeed because the JSON itself is balanced. + # No ; balanced-brace parser still succeeds because + # the JSON itself is balanced. ['{"name":"web_search","arguments":{"query":"x"}}'], ["done"], ], @@ -294,17 +459,11 @@ class TestLoopBasic: assert exec_fn.calls == [("web_search", {"query": "x"})] def test_bad_json_healed_to_query(self): - # Tool call with non-JSON string arguments. With auto_heal_tool_calls - # the string is routed as {"query": ...}. + # Non-JSON string arguments heal to {"query": ...} under auto_heal_tool_calls. loop, exec_fn = _make_loop( turns = [ - # JSON inside the tool call is well-formed; the - # ``arguments`` is a string that is not itself valid - # JSON for ``_coerce_arguments`` to parse, so the - # heal path runs. - [ - '{"name":"web_search","arguments":"hello world"}' - ], + # ``arguments`` is a string _coerce_arguments can't parse, so heal runs. + ['{"name":"web_search","arguments":"hello world"}'], ["ok"], ], exec_results = ["..."], @@ -315,38 +474,203 @@ class TestLoopBasic: class TestLoopBehaviour: - def test_duplicate_tool_call_synthetic_result(self): - # Two identical successful calls in a row: the second is short- - # circuited with a "do not repeat" message and execute_tool is - # called only once. - loop, exec_fn = _make_loop( - turns = [ - [ - '{"name":"web_search","arguments":{"query":"x"}}' - ], - [ - '{"name":"web_search","arguments":{"query":"x"}}' - ], + def test_duplicate_tool_call_internal_noop(self): + captured_messages: list[list[dict]] = [] + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], ["final"], - ], - exec_results = ["search-result-1"], + ] + ) + + def fake_single_turn(messages): + captured_messages.append([dict(message) for message in messages]) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-result-1"]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + ) + + assert exec_fn.calls == [("web_search", {"query": "x"})] + assert [e["tool_call_id"] for e in events if e["type"] == "tool_end"] == ["call_0"] + assert not [ + e + for e in events + if e.get("tool_call_id") == "call_1" and e.get("type") in {"tool_start", "tool_end"} + ] + duplicate_nudges = [ + message + for message in captured_messages[-1] + if message.get("role") == "user" + and "already completed successfully" in message.get("content", "") + ] + assert len(duplicate_nudges) == 1 + + def test_duplicate_tool_call_internal_noop_allows_distinct_followup_tool(self): + captured_messages: list[list[dict]] = [] + captured_tool_names: list[list[str]] = [] + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"python","arguments":{"code":"print(1)"}}'], + ["final"], + ] + ) + + def fake_single_turn(messages, active_tools = None): + captured_messages.append([dict(message) for message in messages]) + captured_tool_names.append( + [ + tool["function"]["name"] + for tool in (active_tools or []) + if tool.get("function", {}).get("name") + ] + ) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-result-1", "python-result"]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 4, + ) + ) + + assert exec_fn.calls == [ + ("web_search", {"query": "x"}), + ("python", {"code": "print(1)"}), + ] + assert [e["tool_call_id"] for e in events if e["type"] == "tool_end"] == [ + "call_0", + "call_2", + ] + assert not [ + e + for e in events + if e.get("tool_call_id") == "call_1" and e.get("type") in {"tool_start", "tool_end"} + ] + duplicate_nudges = [ + message + for message in captured_messages[2] + if message.get("role") == "user" + and "already completed successfully" in message.get("content", "") + ] + assert len(duplicate_nudges) == 1 + assert captured_tool_names[2] == ["web_search", "python"] + + def test_repeated_duplicate_noop_transitions_to_final_attempt(self): + captured_tool_names: list[list[str]] = [] + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["final from first result"], + ] + ) + + def fake_single_turn(messages, active_tools = None): + captured_tool_names.append( + [ + (tool.get("function") or {}).get("name") + for tool in (active_tools or []) + if (tool.get("function") or {}).get("name") + ] + ) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-result"]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 10, + ) + ) + + assert exec_fn.calls == [("web_search", {"query": "x"})] + assert [ + event.get("tool_call_id") for event in events if event.get("type") == "tool_end" + ] == ["call_0"] + assert captured_tool_names[-1] == [] + assert any( + event.get("type") == "content" and "final from first result" in event.get("text", "") + for event in events + ) + + def test_kb_search_capped_per_turn(self): + # Paraphrased KB searches differ by args (dup guard misses them); the + # per-turn cap stops the runaway re-search loop. + n = RAG_MAX_SEARCHES_PER_TURN + queries = [f"paraphrase {i}" for i in range(n + 1)] + turns = [ + [ + '{"name":"search_knowledge_base",' + f'"arguments":{{"query":"{q}"}}}}' + ] + for q in queries + ] + [["final answer"]] + turn_iter = iter(turns) + + def _gen(_messages): + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool([f"chunk-{i}" for i in range(n)]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], + execute_tool = exec_fn, ) events = _collect_events(loop) - # Only one real call. - assert len(exec_fn.calls) == 1 + assert len(exec_fn.calls) == n + assert all(c[0] == "search_knowledge_base" for c in exec_fn.calls) tool_end_events = [e for e in events if e["type"] == "tool_end"] - assert len(tool_end_events) == 2 - assert "do not repeat" in tool_end_events[1]["result"].lower() + assert len(tool_end_events) == n + 1 + assert "do not search again" in tool_end_events[n]["result"].lower() def test_image_sentinel_stripped_from_model_feed(self): - # The tool result has a frontend image sentinel that should be - # stripped before being fed back into the next turn, BUT the - # tool_end event still carries the raw result for the UI. + # The image sentinel is stripped before the next turn, but tool_end still + # carries the raw result for the UI. loop, exec_fn = _make_loop( turns = [ - [ - '{"name":"python","arguments":{"code":"plot()"}}' - ], + ['{"name":"python","arguments":{"code":"plot()"}}'], ["see chart"], ], exec_results = ["chart\n__IMAGES__:/tmp/chart.png"], @@ -379,14 +703,12 @@ class TestLoopBehaviour: auto_heal_tool_calls = True, ) ) - # Model's second turn must not see "__IMAGES__". + # The model's second turn must not see "__IMAGES__". assert len(captured) >= 2 tool_msgs = [m for m in captured[1] if m.get("role") == "tool"] assert tool_msgs, "no tool message reached the model" for tm in tool_msgs: - assert ( - "__IMAGES__" not in tm["content"] - ), f"sentinel leaked to model: {tm['content']!r}" + assert "__IMAGES__" not in tm["content"], f"sentinel leaked to model: {tm['content']!r}" def test_image_sentinel_stripped_with_multiple_markers(self): # Consecutive sentinels: cut at the first, nothing leaks. @@ -416,19 +738,13 @@ class TestLoopBehaviour: tool_msgs = [m for m in captured[1] if m.get("role") == "tool"] assert tool_msgs for tm in tool_msgs: - assert ( - "__IMAGES__" not in tm["content"] - ), f"second sentinel leaked: {tm['content']!r}" - assert ( - tm["content"] == "panel" - ), f"expected payload-only 'panel', got {tm['content']!r}" + assert "__IMAGES__" not in tm["content"], f"second sentinel leaked: {tm['content']!r}" + assert tm["content"] == "panel", f"expected payload-only 'panel', got {tm['content']!r}" def test_tool_execution_error_is_emitted_but_loop_continues(self): loop, exec_fn = _make_loop( turns = [ - [ - '{"name":"web_search","arguments":{"query":"x"}}' - ], + ['{"name":"web_search","arguments":{"query":"x"}}'], ["sorry, that failed"], ], exec_results = ["Error: network unreachable"], @@ -436,16 +752,14 @@ class TestLoopBehaviour: events = _collect_events(loop) tool_end = next(e for e in events if e["type"] == "tool_end") assert tool_end["result"].startswith("Error") - # The loop must still produce a content event after the failure. + # The loop must still emit a content event after the failure. contents = [e for e in events if e["type"] == "content"] assert contents def test_exception_in_executor_does_not_raise(self): loop, exec_fn = _make_loop( turns = [ - [ - '{"name":"web_search","arguments":{"query":"x"}}' - ], + ['{"name":"web_search","arguments":{"query":"x"}}'], ["recovered"], ], exec_results = [RuntimeError("boom")], @@ -459,14 +773,12 @@ class TestLoopControl: def test_cancel_event_breaks_loop(self): cancel = threading.Event() cancel.set() - # Even with a fake stream that emits tool calls, the loop must - # bail before invoking execute_tool when cancel is set. + # With cancel set, the loop bails before invoking execute_tool. exec_fn = FakeExecuteTool([]) events = list( run_safetensors_tool_loop( single_turn = _const_stream( - '{"name":"web_search",' - '"arguments":{"query":"x"}}' + '{"name":"web_search","arguments":{"query":"x"}}' ), messages = [{"role": "user", "content": "hi"}], tools = [], @@ -478,15 +790,13 @@ class TestLoopControl: assert exec_fn.calls == [] def test_max_iterations_caps_loop(self): - # The loop should stop after max_tool_iterations even if the - # model keeps asking for tools, then emit a final-attempt round. + # The loop stops after max_tool_iterations even if the model keeps + # asking for tools, then emits a final-attempt round. loop, exec_fn = _make_loop( turns = [ - # : tool call (executes once) - [ - '{"name":"web_search","arguments":{"query":"a"}}' - ], - # : model gives a final answer when nudged. + # Tool call (executes once). + ['{"name":"web_search","arguments":{"query":"a"}}'], + # Model gives a final answer when nudged. ["here is the final answer"], ], exec_results = ["result"], @@ -494,48 +804,38 @@ class TestLoopControl: ) events = _collect_events(loop) contents = [e for e in events if e["type"] == "content"] - # Final content must include the final answer. + # Final content must contain the final answer. assert contents and "final answer" in contents[-1]["text"] class TestStatusFormatting: def test_status_for_known_tools(self): - # Use the private helper directly to verify status formatting. + # Call the private helper directly to verify status formatting. assert ( - safetensors_agentic._status_for_tool("web_search", {"query": "abc"}) - == "Searching: abc" + safetensors_agentic._status_for_tool("web_search", {"query": "abc"}) == "Searching: abc" ) assert ( - safetensors_agentic._status_for_tool( - "web_search", {"url": "https://www.example.com/x"} - ) + safetensors_agentic._status_for_tool("web_search", {"url": "https://www.example.com/x"}) == "Reading: example.com" ) - assert safetensors_agentic._status_for_tool( - "python", {"code": "x = 1"} - ).startswith("Running Python:") - assert safetensors_agentic._status_for_tool( - "terminal", {"command": "ls"} - ).startswith("Running:") - assert safetensors_agentic._status_for_tool("unknown_tool", {}).startswith( - "Calling:" + assert safetensors_agentic._status_for_tool("python", {"code": "x = 1"}).startswith( + "Running Python:" ) + assert safetensors_agentic._status_for_tool("terminal", {"command": "ls"}).startswith( + "Running:" + ) + assert safetensors_agentic._status_for_tool("unknown_tool", {}).startswith("Calling:") class TestProseMentioningToolCall: def test_assistant_prose_with_literal_tool_call_text_survives(self): - # Regression: if the assistant text legitimately mentions - # ```` as a literal string and the parser finds no - # actual call, the loop must surface the full content instead - # of silently stripping everything past the literal marker. + # Regression: prose that mentions a literal ```` (no real call) + # must surface in full, not be stripped past the marker. loop, exec_fn = _make_loop( turns = [ - # : a real tool call so the loop moves to - # . - [ - '{"name":"web_search","arguments":{"query":"x"}}' - ], - # : prose that mentions the literal text. + # A real tool call so the loop advances a turn. + ['{"name":"web_search","arguments":{"query":"x"}}'], + # Prose that mentions the literal text. ["the docs say means an LLM tool call wrapper"], ], exec_results = ["result"], @@ -549,14 +849,11 @@ class TestProseMentioningToolCall: ), f"prose mentioning should not be truncated; got {final!r}" def test_tool_result_with_tool_call_text_does_not_retrigger(self): - # Tool result text contains the literal ```` string. - # The loop must only parse the MODEL output, not the tool - # result, so we should see exactly one call. + # A literal ```` in the tool result must not re-trigger: the + # loop parses only model output, so exactly one call. loop, exec_fn = _make_loop( turns = [ - [ - '{"name":"web_search","arguments":{"query":"x"}}' - ], + ['{"name":"web_search","arguments":{"query":"x"}}'], ["the docs mention wrappers"], ], exec_results = ["Page text: appears here in the docs"], @@ -572,7 +869,6 @@ class TestChatTemplateHelper: from core.inference.chat_template_helpers import ( apply_chat_template_for_generation, ) - self.apply = apply_chat_template_for_generation class _Tok: @@ -582,7 +878,12 @@ class TestChatTemplateHelper: self.last_kwargs = None def apply_chat_template( - self, messages, *, tokenize = False, add_generation_prompt = True, **kw + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, ): self.call_count += 1 unknown = set(kw) - self.accepted @@ -595,6 +896,7 @@ class TestChatTemplateHelper: tok = self._Tok({"tools", "enable_thinking"}) self.apply(tok, [], tools = [{}], enable_thinking = True) assert tok.call_count == 1 + assert tok.last_kwargs is not None assert "tools" in tok.last_kwargs assert "enable_thinking" in tok.last_kwargs @@ -631,46 +933,66 @@ class TestChatTemplateHelper: class TestGuardrails: def test_disabled_tool_is_not_executed(self): - exec_fn = FakeExecuteTool([]) - loop = run_safetensors_tool_loop( - single_turn = _fake_stream( - [ - '{"name":"terminal","arguments":{"command":"echo bypass"}}' - ] - ), - messages = [{"role": "user", "content": "hi"}], - tools = [{"type": "function", "function": {"name": "web_search"}}], - execute_tool = exec_fn, - max_tool_iterations = 2, - ) - events = _collect_events(loop) - assert exec_fn.calls == [] - tool_ends = [e for e in events if e["type"] == "tool_end"] - assert tool_ends and "not enabled" in tool_ends[0]["result"].lower() + captured_messages: list[list[dict]] = [] - def test_empty_tools_list_does_not_enforce_allowlist(self): - exec_fn = FakeExecuteTool(["OK"]) - loop = run_safetensors_tool_loop( - single_turn = _fake_stream( - [ - '{"name":"python","arguments":{"code":"print(1)"}}' - ] - ), - messages = [{"role": "user", "content": "hi"}], - tools = [], - execute_tool = exec_fn, - max_tool_iterations = 2, + def fake_single_turn(messages): + captured_messages.append([dict(message) for message in messages]) + if len(captured_messages) == 1: + yield '{"name":"terminal","arguments":{"command":"echo bypass"}}' + else: + yield "final" + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + + assert exec_fn.calls == [] + assert not [event for event in events if event.get("type") in {"tool_start", "tool_end"}] + disabled_nudges = [ + message + for message in captured_messages[-1] + if message.get("role") == "user" and "not enabled" in message.get("content", "") + ] + assert len(disabled_nudges) == 1 + + def test_empty_tools_list_means_allow_all_in_core_loop(self): + turns = iter( + [ + ['{"name":"python","arguments":{"code":"print(1)"}}'], + ["done"], + ] + ) + + def fake_single_turn(_messages, active_tools = None): + assert active_tools == [] + acc = "" + for chunk in next(turns): + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["OK"]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) ) - _collect_events(loop) assert exec_fn.calls == [("python", {"code": "print(1)"})] + assert any(event.get("type") == "tool_end" for event in events) def test_max_iterations_zero_executes_no_tools(self): loop, exec_fn = _make_loop( - turns = [ - [ - '{"name":"web_search","arguments":{"query":"x"}}' - ] - ], + turns = [['{"name":"web_search","arguments":{"query":"x"}}']], exec_results = ["OK"], max_tool_iterations = 0, ) @@ -701,9 +1023,7 @@ class TestGuardrails: def test_auto_heal_disabled_still_parses_valid_tool_call(self): loop, exec_fn = _make_loop( turns = [ - [ - '{"name":"web_search","arguments":{"query":"x"}}' - ], + ['{"name":"web_search","arguments":{"query":"x"}}'], ["done"], ], exec_results = ["OK"], @@ -713,50 +1033,125 @@ class TestGuardrails: _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "x"})] + def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self): + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"literal"}}'], + ] + ) + + def fake_single_turn(_messages, active_tools = None): + acc = "" + for chunk in next(turns): + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["OK"]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "show literal"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 1, + auto_heal_tool_calls = False, + ) + ) + assert exec_fn.calls == [("web_search", {"query": "x"})] + assert any( + event.get("type") == "content" and "" in event.get("text", "") + for event in events + ) + + def test_auto_heal_disabled_does_not_repair_unclosed_tool_call(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ], + exec_results = ["OK"], + auto_heal_tool_calls = False, + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + assert any( + event.get("type") == "content" and "" in event.get("text", "") + for event in events + ) + + def test_auto_heal_enabled_strips_unparseable_xml_tool_call(self): + loop, exec_fn = _make_loop( + turns = [["{not valid json}"]], + exec_results = ["OK"], + auto_heal_tool_calls = True, + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + assert not any( + event.get("type") == "content" and "" in event.get("text", "") + for event in events + ) + def test_non_consecutive_duplicate_is_short_circuited(self): loop, exec_fn = _make_loop( turns = [ - [ - '{"name":"web_search","arguments":{"query":"A"}}' - ], - [ - '{"name":"web_search","arguments":{"query":"B"}}' - ], - [ - '{"name":"web_search","arguments":{"query":"A"}}' - ], + ['{"name":"web_search","arguments":{"query":"A"}}'], + ['{"name":"web_search","arguments":{"query":"B"}}'], + ['{"name":"web_search","arguments":{"query":"A"}}'], ["final"], ], exec_results = ["res-A", "res-B"], max_tool_iterations = 4, ) events = _collect_events(loop) - assert exec_fn.calls == [ - ("web_search", {"query": "A"}), - ("web_search", {"query": "B"}), + assert exec_fn.calls == [("web_search", {"query": "A"}), ("web_search", {"query": "B"})] + assert [ + event.get("tool_call_id") for event in events if event.get("type") == "tool_end" + ] == ["call_0", "call_1"] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + + def test_same_turn_duplicate_is_short_circuited(self): + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"A"}}' + '{"name":"web_search","arguments":{"query":"A"}}' + ], + ["final"], + ], + exec_results = ["res-A"], + max_tool_iterations = 2, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "A"})] + assert [ + event.get("tool_call_id") for event in events if event.get("type") == "tool_end" + ] == ["call_0"] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_1" + and event.get("type") in {"tool_start", "tool_end"} ] - tool_ends = [e for e in events if e["type"] == "tool_end"] - assert "already made this exact call" in tool_ends[-1]["result"] def test_coerce_string_args_python_uses_code_key(self): - assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == { - "code": "print(1)" - } + assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"} def test_coerce_string_args_terminal_uses_command_key(self): - assert _coerce_arguments("ls -la", heal = True, tool_name = "terminal") == { - "command": "ls -la" - } + assert _coerce_arguments("ls -la", heal = True, tool_name = "terminal") == {"command": "ls -la"} def test_tool_call_ids_unique_across_loop_iterations(self): loop, _exec = _make_loop( turns = [ - [ - '{"name":"web_search","arguments":{"query":"A"}}' - ], - [ - '{"name":"web_search","arguments":{"query":"B"}}' - ], + ['{"name":"web_search","arguments":{"query":"A"}}'], + ['{"name":"web_search","arguments":{"query":"B"}}'], ["done"], ], exec_results = ["A", "B"], @@ -781,7 +1176,7 @@ class TestGptOssNameDetection: def test_empty_or_none_returns_false(self): assert is_gpt_oss_model_name("") is False - assert is_gpt_oss_model_name(None) is False + assert is_gpt_oss_model_name(cast(str, None)) is False if __name__ == "__main__": diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py index 57007a5f66..24b1da1772 100644 --- a/studio/backend/tests/test_sandbox_tools.py +++ b/studio/backend/tests/test_sandbox_tools.py @@ -88,9 +88,7 @@ class TestTrustedHostAllowlist: _ok(f"import requests; requests.get({url!r})") def test_wikipedia_subdomain_passes(self): - _ok( - 'import urllib.request; urllib.request.urlopen("https://m.en.wikipedia.org/wiki/Foo")' - ) + _ok('import urllib.request; urllib.request.urlopen("https://m.en.wikipedia.org/wiki/Foo")') def test_hf_co_short_form_passes(self): _ok('import requests; requests.get("https://hf.co/unsloth/Qwen3.5-4B-GGUF")') @@ -119,7 +117,7 @@ class TestUntrustedHostBlock: ) def test_dynamic_url_not_statically_blocked(self): - # Static AST cannot resolve runtime URLs; bash blocklist is the fallback. + # Static AST can't resolve runtime URLs; bash blocklist is the fallback. _ok('import requests; url = "https://example.com/"; requests.get(url)') @@ -221,18 +219,12 @@ class TestUploadDenylist: ) def test_plain_post_json_not_blocked(self): - _ok( - "import requests\n" - 'requests.post("https://api.weather.gov/lookup", json={"k": "v"})' - ) + _ok("import requests\n" 'requests.post("https://api.weather.gov/lookup", json={"k": "v"})') class TestSandboxEnvIsolation: - """The sandbox subprocess env is built from a whitelist, not by stripping. - - Confirm every credential-shaped parent var is absent regardless of how the - operator's process is configured. Covers Linux/macOS/WSL/Windows shapes. - """ + """Sandbox env is built from a whitelist, so credential-shaped parent + vars stay absent regardless of operator config (Linux/macOS/WSL/Windows).""" _SECRET_KEYS = ( # HF + ML tooling @@ -318,8 +310,8 @@ class TestSandboxEnvIsolation: def test_term_is_dumb(self, tmp_path): from core.inference.tools import _build_safe_env - # Prevents the sandbox from re-using the operator's TERM (e.g. xterm-256color) - # which could trigger color-escape parsing in downstream tools. + # Avoid re-using the operator's TERM (e.g. xterm-256color) that + # could trigger color-escape parsing in downstream tools. env = _build_safe_env(str(tmp_path)) assert env["TERM"] == "dumb" @@ -345,22 +337,18 @@ class TestSandboxCpuRlimitDefault: class TestMaxBodyDefault: def test_default_is_500_mb(self): - src = (_BACKEND_ROOT / "main.py").read_text() - assert 'UNSLOTH_STUDIO_MAX_BODY_MB", "500"' in src + src = (_BACKEND_ROOT / "utils" / "upload_limits.py").read_text() + assert "DEFAULT_UPLOAD_LIMIT_MB = 500" in src + assert "UNSLOTH_STUDIO_MAX_BODY_MB" in src class TestBashBlocklistPosition: - """The blocklist must fire at command position only. - - Pre-fix the per-token loop fired on any token, so `grep -r curl .` - and `echo source` were rejected. The position-anchored regex plus a - shlex-aware command-position-only token check is sufficient. - """ + """The blocklist must fire at command position only, so args like + `grep -r curl .` and `echo source` are not falsely rejected.""" @staticmethod def _find(): from core.inference.tools import _find_blocked_commands - return _find_blocked_commands # ---- argument-position: must NOT be blocked ---- @@ -371,8 +359,7 @@ class TestBashBlocklistPosition: assert self._find()("echo source the data") == set() def test_cat_with_word_source_allowed(self): - # The 'source' word is an argument to echo; not blocked. - # `echo` itself isn't blocked. Only legit allowed tokens here. + # 'source' is an argument to echo, and echo isn't blocked either. assert self._find()("cat README.md && echo source") == set() assert "source" not in self._find()("cat README.md && echo source") assert "echo" not in self._find()("cat README.md && echo source") @@ -402,14 +389,14 @@ class TestBashBlocklistPosition: assert "wget" in self._find()("cd /tmp && wget https://bad") def test_split_quotes_obfuscation_blocked(self): - # shlex collapses 'r''m' -> 'rm' as a single token at command position. + # shlex collapses 'r''m' -> 'rm' at command position. assert "rm" in self._find()("r''m -rf /") def test_path_prefixed_command_blocked(self): assert "sudo" in self._find()("/usr/bin/sudo whoami") def test_nested_bash_c_blocked(self): - # Recursion into the nested command string still catches command-position curl. + # Recursion into the nested command string catches command-position curl. assert "curl" in self._find()("bash -c 'curl https://x'") def test_subshell_command_blocked(self): @@ -470,9 +457,8 @@ class TestBashBlocklistPosition: class TestHfUploadImportGate: - """HfApi-style upload-method blocking should require an HF import in - scope; otherwise paramiko / boto3 / internal SDKs with the same - method names hit a false positive.""" + """Upload-method blocking requires an HF import in scope, so paramiko / + boto3 / internal SDKs with the same method names don't false-positive.""" def test_paramiko_upload_file_allowed_without_hf_import(self): _ok("import paramiko; sftp=None; sftp.upload_file('a','b')") @@ -481,14 +467,14 @@ class TestHfUploadImportGate: _ok("client=None; client.create_commit(Repo='x')") def test_hf_api_upload_safe_path_allowed(self): - # Sandbox-local relative path -- the call shape we want to permit. + # Sandbox-local relative path -- the permitted call shape. _ok("from huggingface_hub import HfApi; HfApi().upload_file('a','b','c')") def test_hf_upload_file_fq_safe_path_allowed(self): _ok("import huggingface_hub; huggingface_hub.upload_file('a','b','c')") def test_dynamic_builtin_import_safe_path_allowed(self): - # `__import__('huggingface_hub')` puts HF in scope; relative-literal path is safe. + # `__import__('huggingface_hub')` puts HF in scope; relative literal is safe. _ok("hf=__import__('huggingface_hub'); hf.HfApi().upload_file('a','b','c')") def test_dynamic_importlib_safe_path_allowed(self): @@ -504,8 +490,8 @@ class TestHfUploadImportGate: ) def test_hf_bare_name_upload_safe_path_allowed(self): - # `from huggingface_hub import upload_file` then bare `upload_file(...)` - # with a sandbox-local relative-path literal is allowed. + # Bare `upload_file(...)` (imported from huggingface_hub) with a + # sandbox-local relative-path literal is allowed. _ok( "from huggingface_hub import upload_file;" " upload_file(path_or_fileobj='x', path_in_repo='x', repo_id='r')" @@ -524,15 +510,14 @@ class TestHfUploadImportGate: ) def test_bare_name_upload_file_without_hf_import_allowed(self): - # No HF import -- local helper named upload_file should pass. - _ok("def upload_file(*a, **k):\n pass\n" "upload_file('x', 'y', 'z')") + # No HF import -- local helper named upload_file passes. + _ok("def upload_file(*a, **k):\n pass\nupload_file('x', 'y', 'z')") class TestHfUploadSandboxLocalPaths: - """The HF upload gate must only allow uploads of files that already live in - the sandbox workdir. Absolute paths, `..` traversal, home expansion, and - Windows drive letters are rejected because the LLM can use them to lift - secrets from outside the sandbox.""" + """HF upload gate allows only files in the sandbox workdir. Absolute paths, + `..` traversal, home expansion, and Windows drives are rejected (they could + lift secrets from outside the sandbox).""" def test_relative_literal_allowed(self): _ok( @@ -626,8 +611,8 @@ class TestHfUploadSandboxLocalPaths: ) def test_dynamic_variable_path_blocked(self): - # A non-literal expression could resolve to any path at runtime; - # the static checker cannot prove safety, so block. + # A non-literal expr could resolve to any path at runtime; the + # static checker can't prove safety, so block. _blocked( "import huggingface_hub, os\n" "p = os.path.join('outputs', 'x.bin')\n" @@ -679,11 +664,9 @@ class TestHfUploadSandboxLocalPaths: class TestHfUploadEnvAndSecretLeakBlock: - """The HF upload gate must reject any positional / keyword arg sourced from - `os.environ` / `os.getenv` / subprocess env reads. Even though - `_build_safe_env` strips HF_TOKEN/WANDB/AWS upfront for the sandbox shell, - a Python script can still reach the parent process env if it bypasses the - safe-env wrapper at the source -- so block statically.""" + """HF upload gate rejects any arg sourced from os.environ / os.getenv / + subprocess env reads, since a script can reach the parent env directly + despite the safe-env shell wrapper.""" def test_path_from_os_environ_subscript_blocked(self): _blocked( @@ -761,7 +744,7 @@ class TestHfUploadEnvAndSecretLeakBlock: ) def test_env_dict_unpacked_via_environ_attr_blocked(self): - # `os.environ` as a bare reference (passed somewhere it gets serialized). + # Bare `os.environ` reference (passed somewhere it gets serialized). _blocked( "import huggingface_hub, os\n" "huggingface_hub.upload_file(path_or_fileobj=str(os.environ)," @@ -770,8 +753,8 @@ class TestHfUploadEnvAndSecretLeakBlock: ) def test_repo_id_from_env_also_blocked(self): - # Even non-path args must not source env vars -- an attacker could - # encode secrets in repo_id or path_in_repo. + # Non-path args must not source env vars either -- an attacker + # could encode secrets in repo_id or path_in_repo. _blocked( "import huggingface_hub, os\n" 'huggingface_hub.upload_file(path_or_fileobj="x.bin",' diff --git a/studio/backend/tests/test_server_disk_logging.py b/studio/backend/tests/test_server_disk_logging.py new file mode 100644 index 0000000000..05d03d869c --- /dev/null +++ b/studio/backend/tests/test_server_disk_logging.py @@ -0,0 +1,99 @@ +# 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 the server session log + native-crash capture in run.py. + +Field regression: Studio "terminates without a warning" -- a native crash in +the GPU runtime kills the process with no Python traceback, and a desktop- +shortcut console closes before anything can be read. The server must tee its +console output to disk and aim faulthandler at the same file so even hard +crashes leave evidence. +""" + +from __future__ import annotations + +import io +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import run as run_mod # noqa: E402 + + +class TestTeeStream: + def test_writes_reach_both_and_return_original(self): + console, log = io.StringIO(), io.StringIO() + tee = run_mod._TeeStream(console, log) + n = tee.write("hello") + assert console.getvalue() == "hello" == log.getvalue() + assert n == 5 # delegate's return value, console contract unchanged + + def test_log_failure_never_breaks_console(self): + class Broken: + def write(self, data): + raise OSError("disk full") + + def flush(self): + raise OSError("disk full") + + console = io.StringIO() + tee = run_mod._TeeStream(console, Broken()) + assert tee.write("still works") == len("still works") + tee.flush() # must not raise + assert console.getvalue() == "still works" + + def test_attribute_proxy(self): + console, log = io.StringIO(), io.StringIO() + tee = run_mod._TeeStream(console, log) + # isatty / encoding probes must see the original stream's answers. + assert tee.isatty() == console.isatty() + + +class TestSetupServerDiskLogging: + def test_opt_out_env(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_NO_FILE_LOG", "1") + assert run_mod._setup_server_disk_logging() is None + + def test_creates_log_and_enables_faulthandler(self, monkeypatch, tmp_path): + import faulthandler + + monkeypatch.delenv("UNSLOTH_STUDIO_NO_FILE_LOG", raising = False) + monkeypatch.delenv("PYTHONFAULTHANDLER", raising = False) + # Both resolution paths (utils.paths.studio_root and the env + # fallback) honor UNSLOTH_STUDIO_HOME, so this redirects the log dir. + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + orig_out, orig_err = sys.stdout, sys.stderr + was_enabled = faulthandler.is_enabled() + try: + log_path = run_mod._setup_server_disk_logging() + assert log_path is not None + assert Path(log_path).is_file() + assert "logs" in str(log_path) + # faulthandler armed at the file; children inherit the env switch. + assert faulthandler.is_enabled() + import os + + assert os.environ.get("PYTHONFAULTHANDLER") == "1" + print("tee-capture-marker") + sys.stdout.flush() + assert "tee-capture-marker" in Path(log_path).read_text( + encoding = "utf-8", errors = "replace" + ) + finally: + sys.stdout, sys.stderr = orig_out, orig_err + if not was_enabled: + faulthandler.disable() + + def test_run_server_wires_logging_before_main_import(self): + src = (Path(_BACKEND_DIR) / "run.py").read_text(encoding = "utf-8") + call_idx = src.index("_setup_server_disk_logging()", src.index("def run_server")) + main_import_idx = src.index("from main import app", src.index("def run_server")) + assert call_idx < main_import_idx, ( + "disk logging must be armed before importing main so import-time " + "failures leave evidence on disk" + ) diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py index 521c99e126..27f695b744 100644 --- a/studio/backend/tests/test_studio_api.py +++ b/studio/backend/tests/test_studio_api.py @@ -4,9 +4,9 @@ """ End-to-end tests for Unsloth Studio's HTTP API surface. -Covers the OpenAI-compatible and Anthropic-compatible endpoints exposed -by the server that ``unsloth studio run`` boots, plus API key -authentication and the CLI's ``--help`` output: +Covers the OpenAI- and Anthropic-compatible endpoints exposed by the +server that ``unsloth studio run`` boots, plus API key authentication and +the CLI's ``--help`` output: 1. curl -- basic chat completions (non-streaming) 2. curl -- streaming chat completions @@ -38,14 +38,14 @@ Usage: export UNSLOTH_E2E_API_KEY=sk-unsloth-... # from the server banner pytest tests/test_studio_api.py -v - # Pytest mode, fixture-managed server — pytest launches and tears - # down the server itself. One-shot verification, CI-friendly. + # Pytest mode, fixture-managed server — pytest launches and tears down + # the server itself. One-shot verification, CI-friendly. pytest tests/test_studio_api.py -v \\ --unsloth-model unsloth/Qwen3-1.7B-GGUF \\ --unsloth-gguf-variant UD-Q4_K_XL -The ``base_url`` / ``api_key`` parameters on the test functions resolve -via the ``studio_server`` session fixture in ``conftest.py``. +The ``base_url`` / ``api_key`` parameters on the test functions resolve via +the ``studio_server`` session fixture in ``conftest.py``. Requires a GPU and ~2 GB of disk for the GGUF download. """ @@ -65,21 +65,17 @@ import urllib.request from pathlib import Path -# ── Configuration ──────────────────────────────────────────────────── +# Configuration DEFAULT_MODEL = "unsloth/Qwen3-1.7B-GGUF" DEFAULT_VARIANT = "UD-Q4_K_XL" PORT = 18222 # high port unlikely to collide HOST = "127.0.0.1" -STARTUP_TIMEOUT = 120 # seconds to wait for banner -LOG_FILE = ( - Path(__file__).resolve().parent.parent.parent.parent - / "temp" - / "test_studio_api.log" -) +STARTUP_TIMEOUT = 120 # seconds +LOG_FILE = Path(__file__).resolve().parent.parent.parent.parent / "temp" / "test_studio_api.log" -# ── Helpers ────────────────────────────────────────────────────────── +# Helpers def _http( @@ -129,7 +125,7 @@ def _stream_http( return exc.code, [] -# ── Test functions ─────────────────────────────────────────────────── +# Test functions def test_help_output(): @@ -219,9 +215,7 @@ def test_openai_sdk(base_url: str, api_key: str): client = OpenAI(base_url = f"{base_url}/v1", api_key = api_key) response = client.chat.completions.create( model = "current", - messages = [ - {"role": "user", "content": "What is 2+2? Answer with just the number."} - ], + messages = [{"role": "user", "content": "What is 2+2? Answer with just the number."}], stream = True, ) content_parts = [] @@ -239,10 +233,10 @@ def test_openai_sdk(base_url: str, api_key: str): def test_curl_with_tools(base_url: str, api_key: str): """Example 4: chat completion with tool calling enabled. - Note: when ``enable_tools`` is set the server always returns SSE - streaming regardless of the ``stream`` flag, so we parse SSE chunks. - The model may or may not produce visible content -- tool orchestration - can intercept the response -- so we only assert the endpoint succeeds. + When ``enable_tools`` is set the server always returns SSE streaming + regardless of the ``stream`` flag, so we parse SSE chunks. The model may + not produce visible content (tool orchestration can intercept the + response), so we only assert the endpoint succeeds. """ status, chunks = _stream_http( f"{base_url}/v1/chat/completions", @@ -271,19 +265,13 @@ def test_curl_with_tools(base_url: str, api_key: str): print(f" PASS curl with tools: {len(chunks)} chunks, {len(full)} chars content") -# ── Standard OpenAI function-calling pass-through tests ───────────── +# Standard OpenAI function-calling pass-through tests. # -# Regression coverage for unslothai/unsloth#4999: Studio's -# /v1/chat/completions used to silently strip standard OpenAI `tools` -# and `tool_choice` fields, so clients (opencode, Claude Code, Cursor, -# Continue, ...) could never get structured tool_calls back. These -# tests exercise the client-side pass-through path that forwards those -# fields to llama-server verbatim. -# -# They require a tool-capable GGUF (``supports_tools=True`` — e.g. -# Qwen3, Qwen2.5-Coder, Llama-3.1-Instruct). The default test model -# ``unsloth/Qwen3-1.7B-GGUF`` advertises tool support via its chat -# template metadata. +# Regression coverage for unslothai/unsloth#4999: /v1/chat/completions used +# to strip standard OpenAI `tools`/`tool_choice`, so clients never got +# structured tool_calls back. These exercise the pass-through that forwards +# those fields to llama-server verbatim. Require a tool-capable GGUF +# (supports_tools=True); the default unsloth/Qwen3-1.7B-GGUF qualifies. _WEATHER_TOOL = { "type": "function", @@ -307,9 +295,9 @@ _WEATHER_TOOL = { def _collect_streamed_tool_calls(chunks: list[dict]) -> list[dict]: """Reassemble OpenAI streaming delta.tool_calls into full tool calls. - OpenAI streams partial tool calls across chunks — the first chunk for - a given index carries ``id`` + ``function.name``, and subsequent - chunks append fragments to ``function.arguments``. + OpenAI streams partial tool calls across chunks — the first chunk for a + given index carries ``id`` + ``function.name``, and later chunks append + fragments to ``function.arguments``. """ by_index: dict[int, dict] = {} for c in chunks: @@ -352,8 +340,8 @@ def _final_finish_reason(chunks: list[dict]) -> str | None: def test_openai_tools_nonstream(base_url: str, api_key: str): """Standard OpenAI function calling, non-streaming, tool_choice='required'. - Regression: before the fix, Studio silently stripped `tools` and the - model returned plain text with finish_reason='stop'. After the fix, + Regression: before the fix, Studio stripped `tools` and the model + returned plain text with finish_reason='stop'. After the fix, llama-server's response is forwarded verbatim so the client sees finish_reason='tool_calls' with a structured tool_calls array and non-zero usage.prompt_tokens. @@ -390,9 +378,7 @@ def test_openai_tools_nonstream(base_url: str, api_key: str): assert "city" in parsed, f"Tool call missing required 'city' arg: {parsed}" # Usage must be non-zero (was 0 before the fix) usage = data.get("usage") or {} - assert ( - usage.get("prompt_tokens", 0) > 0 - ), f"Expected non-zero prompt_tokens; got {usage}" + assert usage.get("prompt_tokens", 0) > 0, f"Expected non-zero prompt_tokens; got {usage}" assert data.get("id"), "Missing response id" print( f" PASS openai tools non-stream: " @@ -417,8 +403,7 @@ def test_openai_tools_stream(base_url: str, api_key: str): assert status == 200, f"Expected 200, got {status}" assert len(chunks) > 0, "No SSE chunks received" assert _final_finish_reason(chunks) == "tool_calls", ( - f"Expected final finish_reason='tool_calls', got " - f"{_final_finish_reason(chunks)!r}" + f"Expected final finish_reason='tool_calls', got " f"{_final_finish_reason(chunks)!r}" ) assembled = _collect_streamed_tool_calls(chunks) assert len(assembled) >= 1, "No tool_calls reassembled from stream" @@ -437,9 +422,9 @@ def test_openai_tools_multiturn(base_url: str, api_key: str): messages and assistant messages carrying tool_calls are accepted. Regression: before the fix, ChatMessage.role was restricted to - {system,user,assistant} and rejected role='tool' at the Pydantic - validation stage. This test sends a full round trip so the model - receives the simulated tool result and responds with final text. + {system,user,assistant} and rejected role='tool' at Pydantic + validation. This test sends a full round trip so the model receives the + simulated tool result and responds with final text. """ status, text = _http( "POST", @@ -476,7 +461,7 @@ def test_openai_tools_multiturn(base_url: str, api_key: str): assert status == 200, f"Expected 200, got {status}: {text[:500]}" data = json.loads(text) msg = data["choices"][0]["message"] - # The model should respond with text now that it has the tool result + # The model should respond with text now it has the tool result content = msg.get("content") or "" assert len(content) > 0 or msg.get( "tool_calls" @@ -501,8 +486,7 @@ def test_openai_sdk_tool_calling(base_url: str, api_key: str): stream = False, ) assert resp.choices[0].finish_reason == "tool_calls", ( - f"Expected finish_reason='tool_calls', got " - f"{resp.choices[0].finish_reason!r}" + f"Expected finish_reason='tool_calls', got " f"{resp.choices[0].finish_reason!r}" ) tool_calls = resp.choices[0].message.tool_calls assert tool_calls and len(tool_calls) >= 1, "No tool_calls from SDK" @@ -510,9 +494,7 @@ def test_openai_sdk_tool_calling(base_url: str, api_key: str): assert tc.function.name == "get_weather" parsed = json.loads(tc.function.arguments) assert "city" in parsed - print( - f" PASS openai SDK tool calling: " f"tool={tc.function.name}, args={parsed}" - ) + print(f" PASS openai SDK tool calling: " f"tool={tc.function.name}, args={parsed}") def test_invalid_key_rejected(base_url: str): @@ -544,7 +526,7 @@ def test_no_key_rejected(base_url: str): print(f" PASS no API key rejected ({status})") -# ── Anthropic SSE helper ───────────────────────────────────────────── +# Anthropic SSE helper def _stream_anthropic_http( @@ -592,7 +574,7 @@ def _collect_anthropic_text(events: list[tuple[str, dict]]) -> str: return "".join(parts) -# ── Anthropic /v1/messages test functions ──────────────────────────── +# Anthropic /v1/messages test functions def test_anthropic_basic(base_url: str, api_key: str): @@ -655,9 +637,7 @@ def test_anthropic_sdk(base_url: str, api_key: str): message = client.messages.create( model = "default", max_tokens = 100, - messages = [ - {"role": "user", "content": "What is 2+2? Answer with just the number."} - ], + messages = [{"role": "user", "content": "What is 2+2? Answer with just the number."}], ) assert message.role == "assistant" assert len(message.content) > 0, "Empty content" @@ -708,18 +688,16 @@ def test_anthropic_with_tools(base_url: str, api_key: str): assert "message_stop" in event_types, "Missing message_stop" full = _collect_anthropic_text(events) - print( - f" PASS anthropic with tools: {len(events)} events, {len(full)} chars content" - ) + print(f" PASS anthropic with tools: {len(events)} events, {len(full)} chars content") def test_anthropic_tool_choice_any(base_url: str, api_key: str): """Anthropic Messages API: ``tool_choice: {"type": "any"}`` must be honored (forwarded as OpenAI ``tool_choice: "required"`` to llama-server). Regression for the secondary fix bundled with #4999 — - previously this field was accepted on the request model but silently - dropped with a warning log, so the model was free to answer from - memory instead of using the tool. + previously this field was accepted on the request model but dropped with + a warning log, so the model could answer from memory instead of using + the tool. """ status, events = _stream_anthropic_http( f"{base_url}/v1/messages", @@ -727,7 +705,7 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str): "model": "default", "max_tokens": 256, "messages": [ - # A question the model could easily answer from memory if + # A question the model could answer from memory if # tool_choice were not enforced. { "role": "user", @@ -756,7 +734,7 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str): assert status == 200, f"Expected 200, got {status}" assert len(events) > 0, "No SSE events received" - # With tool_choice=any, stop_reason must be tool_use (not end_turn) + # With tool_choice=any, stop_reason must be tool_use, not end_turn stop_reason = None for etype, data in events: if etype == "message_delta": @@ -770,8 +748,7 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str): tool_use_starts = [ e for e in events - if e[0] == "content_block_start" - and e[1].get("content_block", {}).get("type") == "tool_use" + if e[0] == "content_block_start" and e[1].get("content_block", {}).get("type") == "tool_use" ] assert len(tool_use_starts) >= 1, "No tool_use content block emitted" print( @@ -780,7 +757,7 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str): ) -# ── Server lifecycle ───────────────────────────────────────────────── +# Server lifecycle def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, str]: @@ -821,9 +798,7 @@ def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, st if proc.poll() is not None: log_fh.flush() log_text = LOG_FILE.read_text() - raise RuntimeError( - f"Server exited early (code {proc.returncode}):\n{log_text[-2000:]}" - ) + raise RuntimeError(f"Server exited early (code {proc.returncode}):\n{log_text[-2000:]}") log_text = LOG_FILE.read_text() m = re.search(r"API Key:\s+(sk-unsloth-[a-f0-9]+)", log_text) if m: @@ -833,9 +808,7 @@ def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, st if not api_key: log_text = LOG_FILE.read_text() _kill_server(proc) - raise RuntimeError( - f"Timed out waiting for API key in server output:\n{log_text[-2000:]}" - ) + raise RuntimeError(f"Timed out waiting for API key in server output:\n{log_text[-2000:]}") # Wait a moment for the model to be fully loaded time.sleep(2) @@ -858,13 +831,11 @@ def _kill_server(proc: subprocess.Popen): proc.wait(timeout = 5) -# ── Main ───────────────────────────────────────────────────────────── +# Main def main(): - parser = argparse.ArgumentParser( - description = "End-to-end tests for unsloth studio run" - ) + parser = argparse.ArgumentParser(description = "End-to-end tests for unsloth studio run") parser.add_argument( "--model", default = DEFAULT_MODEL, @@ -893,14 +864,12 @@ def main(): failed += 1 print(f" ERROR {fn.__name__}: {type(exc).__name__}: {exc}") - # ── 1. Test --help (no server needed) ──────────────────────────── + # 1. --help (no server needed) print("\n[1/16] Testing --help output") run_test(test_help_output) - # ── 2-16. Start server and run API tests ───────────────────────── - print( - f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}..." - ) + # 2-16. Start server and run API tests + print(f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}...") proc = None try: proc, api_key = _start_server(args.model, args.gguf_variant) @@ -954,14 +923,14 @@ def main(): except RuntimeError as exc: print(f"\nFATAL: Server failed to start: {exc}") - failed += 16 # count remaining tests as failed + failed += 16 # remaining tests count as failed finally: if proc: print("\nStopping server...") _kill_server(proc) print("Server stopped.") - # ── Summary ────────────────────────────────────────────────────── + # Summary total = passed + failed print(f"\n{'=' * 40}") print(f"Results: {passed}/{total} passed, {failed} failed") diff --git a/studio/backend/tests/test_studio_train_validation.py b/studio/backend/tests/test_studio_train_validation.py index 7ffa9bb384..0ffecb3ce4 100644 --- a/studio/backend/tests/test_studio_train_validation.py +++ b/studio/backend/tests/test_studio_train_validation.py @@ -18,11 +18,13 @@ from models.training import ( _MAX_LORA_ALPHA, _MAX_LORA_R, _MAX_SEQ_LENGTH, + _MAX_VISION_IMAGE_SIZE, + _MIN_VISION_IMAGE_SIZE, ) def _check_field(field_name: str, value): - """Run the field validator without constructing a full TrainingStartRequest.""" + """Run the field validator without building a full TrainingStartRequest.""" from models.training import TrainingStartRequest schema_field = TrainingStartRequest.model_fields[field_name] @@ -62,6 +64,50 @@ class TestBatchSizeCap: _check_field("batch_size", 0) +class TestVisionImageSizeCap: + def test_none_accepts_model_default(self): + _check_field("vision_image_size", None) + + @pytest.mark.parametrize( + "value", + [_MIN_VISION_IMAGE_SIZE, 640, 1000, _MAX_VISION_IMAGE_SIZE], + ) + def test_in_range_accepts(self, value): + _check_field("vision_image_size", value) + assert _MIN_VISION_IMAGE_SIZE == 256 + assert _MAX_VISION_IMAGE_SIZE == 2048 + + @pytest.mark.parametrize( + "value", + [_MIN_VISION_IMAGE_SIZE - 1, _MAX_VISION_IMAGE_SIZE + 1, 640.5, True], + ) + def test_invalid_rejects(self, value): + with pytest.raises(ValidationError): + _check_field("vision_image_size", value) + + @pytest.mark.parametrize("value", [True, False]) + def test_bool_error_says_integer_not_range(self, value): + # Regression guard: bools say "integer or null", not "in [256, 2048]". + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + + @pytest.mark.parametrize("value", ["++512", "--256", "+-+512", "+", "-"]) + def test_multi_sign_string_says_integer_not_raw(self, value): + # Regression guard: multi-sign strings say "integer or null", not int()'s raw message. + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + assert "invalid literal" not in str(exc.value) + + @pytest.mark.parametrize("value", ["５１２", "٥١٢", "१०२४"]) + def test_unicode_digit_string_rejected(self, value): + # Reject non-ASCII (full-width/Arabic-Indic/Devanagari) digits. + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + + class TestLoraRCap: def test_at_cap_accepts(self): _check_field("lora_r", _MAX_LORA_R) diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py new file mode 100644 index 0000000000..8ff41342d7 --- /dev/null +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Strict-mode (Auto-Heal disabled) tool-call parsing. + +With ``allow_incomplete=False`` the parser must accept a well-formed +``...`` call even when the model appends prose +after the closing tag -- matching the JSON-style ``...`` path, +which already tolerates trailing text -- while still rejecting genuinely +truncated calls that never close. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.tool_call_parser import parse_tool_calls_from_text + + +def _only(text: str) -> dict: + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1, f"expected exactly one call, got {len(calls)}: {calls!r}" + fn = calls[0]["function"] + return {"name": fn["name"], "arguments": json.loads(fn["arguments"])} + + +class TestFunctionStyleTrailingText: + def test_closed_function_with_trailing_prose_is_accepted(self): + text = ( + "weather london" + " Let me check that for you." + ) + call = _only(text) + assert call == {"name": "web_search", "arguments": {"query": "weather london"}} + + def test_closed_function_with_trailing_whitespace_is_accepted(self): + text = "cats \n\n" + call = _only(text) + assert call == {"name": "web_search", "arguments": {"query": "cats"}} + + def test_closed_function_without_trailing_text_still_parses(self): + text = "cats" + call = _only(text) + assert call == {"name": "web_search", "arguments": {"query": "cats"}} + + def test_multi_param_with_trailing_prose(self): + text = ( + "ls -la" + "home running it now" + ) + call = _only(text) + assert call == { + "name": "terminal", + "arguments": {"command": "ls -la", "workdir": "home"}, + } + + def test_code_value_containing_literal_close_tag_is_preserved(self): + # The real closing is the last one; the literal inside + # the code argument must survive (rfind, not the first match). + text = ( + "" + 'print("")' + " all done" + ) + call = _only(text) + assert call == {"name": "python", "arguments": {"code": 'print("")'}} + + def test_incomplete_function_without_close_is_still_rejected(self): + text = "weather london" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_param_without_close_tag_is_rejected_in_strict_mode(self): + # Closing present, but the single parameter never closes. + text = "weather london" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + +class TestParityWithJsonStyle: + def test_json_tool_call_with_trailing_prose_is_accepted(self): + text = ( + '{"name":"web_search","arguments":{"query":"weather london"}}' + " Let me check that for you." + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + def test_function_and_json_styles_agree_on_trailing_text(self): + q = "weather london" + func = parse_tool_calls_from_text( + f"{q} trailing", + allow_incomplete = False, + ) + js = parse_tool_calls_from_text( + f'{{"name":"web_search","arguments":{{"query":"{q}"}}}} trailing', + allow_incomplete = False, + ) + assert len(func) == len(js) == 1 + assert json.loads(func[0]["function"]["arguments"]) == {"query": q} + assert json.loads(js[0]["function"]["arguments"]) == {"query": q} + + +class TestHealingPathUnaffected: + def test_auto_heal_still_repairs_unclosed_function(self): + text = "cats" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" diff --git a/studio/backend/tests/test_tool_loop_controller.py b/studio/backend/tests/test_tool_loop_controller.py new file mode 100644 index 0000000000..dea5de6d6e --- /dev/null +++ b/studio/backend/tests/test_tool_loop_controller.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.tool_loop_controller import ( + ToolLoopController, + canonical_tool_call_key, + coerce_tool_arguments, + status_for_tool, + strip_result_for_model, + tool_event_provenance, +) + + +def _tool(name: str) -> dict: + return {"type": "function", "function": {"name": name}} + + +def _call( + name: str, + args, + call_id: str = "call_0", +) -> dict: + return { + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args) if isinstance(args, dict) else args, + }, + } + + +def test_canonical_tool_call_key_sorts_arguments(): + a = canonical_tool_call_key("web_search", {"query": "gpu", "limit": 5}) + b = canonical_tool_call_key("web_search", {"limit": 5, "query": "gpu"}) + c = canonical_tool_call_key("python", {"limit": 5, "query": "gpu"}) + + assert a == b + assert a != c + assert a == 'web_search:{"limit":5,"query":"gpu"}' + + +def test_coerce_tool_arguments_parses_json_and_heals_raw_strings(): + parsed = coerce_tool_arguments('{"query":"gpu prices"}', heal = True) + healed = coerce_tool_arguments("print(1)", heal = True, tool_name = "python") + raw = coerce_tool_arguments("not-json", heal = False, tool_name = "python") + + assert parsed.arguments == {"query": "gpu prices"} + assert not parsed.healed + assert healed.arguments == {"code": "print(1)"} + assert healed.healed + assert raw.arguments == {"raw": "not-json"} + assert not raw.healed + + +def test_status_and_provenance_match_local_event_conventions(): + assert status_for_tool("web_search", {"query": "gpus"}) == "Searching: gpus" + assert ( + status_for_tool("web_search", {"url": "https://www.example.com/a"}) + == "Reading: example.com" + ) + assert status_for_tool("python", {"code": "print(1)\nprint(2)"}) == "Running Python: print(1)" + assert tool_event_provenance(healed = True, forced = False, provisional = None) == { + "source": "local", + "healed": True, + } + + +def test_prepare_execute_builds_visible_events_and_model_tool_message(): + controller = ToolLoopController(tools = [_tool("web_search")]) + decision = controller.prepare_call(_call("web_search", {"query": "gpu prices"})) + + assert decision.should_execute + assert decision.emit_visible_events + assert decision.status_text == "Searching: gpu prices" + assert decision.tool_start_payload()["arguments"] == {"query": "gpu prices"} + assert decision.tool_start_event()["type"] == "tool_start" + assert decision.as_assistant_tool_call()["function"]["arguments"] == '{"query":"gpu prices"}' + + completion = controller.record_result(decision, "Search result\n__IMAGES__:{...}") + + assert completion.tool_end_payload()["result"] == "Search result\n__IMAGES__:{...}" + assert completion.tool_end_event()["type"] == "tool_end" + assert completion.tool_message() == { + "role": "tool", + "name": "web_search", + "content": "Search result", + "tool_call_id": "call_0", + } + + +def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools(): + controller = ToolLoopController(tools = [_tool("web_search"), _tool("python")]) + first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_a")) + controller.record_result(first, "ok") + + duplicate = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_b")) + completion = controller.record_noop(duplicate) + + assert duplicate.action == "duplicate" + assert not duplicate.should_execute + assert not duplicate.emit_visible_events + duplicate_nudge = completion.model_message()["content"] + assert "already completed successfully" in duplicate_nudge + assert "different enabled tool" in duplicate_nudge + assert completion.model_message()["role"] == "user" + assert not controller.force_final_answer + assert [tool["function"]["name"] for tool in controller.active_tools()] == [ + "web_search", + "python", + ] + + +def test_repeated_successful_duplicate_becomes_terminal_after_one_recovery_nudge(): + controller = ToolLoopController(tools = [_tool("web_search"), _tool("python")]) + first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_a")) + controller.record_result(first, "ok") + + duplicate_one = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_b")) + completion_one = controller.record_noop(duplicate_one) + + assert duplicate_one.action == "duplicate" + assert "already completed successfully" in completion_one.model_message()["content"] + assert not controller.force_final_answer + assert [tool["function"]["name"] for tool in controller.active_tools()] == [ + "web_search", + "python", + ] + + duplicate_two = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_c")) + completion_two = controller.record_noop(duplicate_two) + + assert duplicate_two.action == "duplicate" + assert "already completed successfully" in completion_two.model_message()["content"] + assert controller.force_final_answer + assert controller.active_tools() == [] + + +def test_failed_call_does_not_block_retry(): + controller = ToolLoopController(tools = [_tool("web_search")]) + first = controller.prepare_call(_call("web_search", {"query": "gpu prices"})) + controller.record_result(first, "Error: temporary failure") + + retry = controller.prepare_call(_call("web_search", {"query": "gpu prices"})) + + assert retry.should_execute + assert retry.action == "execute" + + +def test_empty_enabled_tool_list_blocks_all_tool_calls(): + controller = ToolLoopController(tools = []) + decision = controller.prepare_call(_call("web_search", {"query": "gpu prices"})) + completion = controller.record_noop(decision) + + assert decision.action == "disabled" + assert not decision.emit_visible_events + assert completion.model_message()["role"] == "user" + assert "not enabled" in completion.model_message()["content"] + assert controller.force_final_answer + assert controller.active_tools() == [] + + +def test_disabled_tool_is_internal_noop_not_visible_tool_error(): + controller = ToolLoopController(tools = [_tool("web_search")]) + decision = controller.prepare_call(_call("python", {"code": "print(1)"})) + completion = controller.record_noop(decision) + + assert decision.action == "disabled" + assert not decision.emit_visible_events + assert completion.model_message()["role"] == "user" + assert "not enabled" in completion.model_message()["content"] + assert controller.force_final_answer + assert controller.active_tools() == [] + + +def test_render_html_success_filters_active_tools_and_repeat_is_internal(): + controller = ToolLoopController(tools = [_tool("render_html"), _tool("web_search")]) + assert [t["function"]["name"] for t in controller.active_tools()] == [ + "render_html", + "web_search", + ] + + first = controller.prepare_call(_call("render_html", {"code": ""}, "call_html_1")) + controller.record_result(first, "Rendered HTML artifact: Demo") + + assert [t["function"]["name"] for t in controller.active_tools()] == ["web_search"] + + repeat = controller.prepare_call(_call("render_html", {"code": ""}, "call_html_2")) + completion = controller.record_noop(repeat) + + assert repeat.action == "render_html_repeat" + assert not repeat.emit_visible_events + assert completion.model_message()["role"] == "user" + assert "Do not call render_html again" in completion.model_message()["content"] + assert controller.force_final_answer + assert controller.active_tools() == [] + + +def test_strip_result_for_model_removes_frontend_image_sentinel(): + assert strip_result_for_model('text\n__IMAGES__:{"paths":[]}') == "text" + assert strip_result_for_model("text __IMAGES__:payload") == "text" + assert strip_result_for_model("plain text") == "plain text" diff --git a/studio/backend/tests/test_tool_message_empty_content.py b/studio/backend/tests/test_tool_message_empty_content.py new file mode 100644 index 0000000000..d63b16ce80 --- /dev/null +++ b/studio/backend/tests/test_tool_message_empty_content.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Empty ``role="tool"`` content must be accepted on the OpenAI-compat surface. + +Agentic clients send ``content: ""`` when a command produced no output; +OpenAI and llama-server both accept it. Studio used to 400, which standard +clients treat as non-retryable and kill the session. The validator must +normalize empty/missing tool content to ``""`` instead of raising. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from models.inference import ChatMessage + + +def test_tool_message_empty_string_content_is_accepted(): + msg = ChatMessage(role = "tool", content = "", tool_call_id = "call_1") + assert msg.content == "" + + +def test_tool_message_none_content_normalizes_to_empty_string(): + msg = ChatMessage(role = "tool", content = None, tool_call_id = "call_1") + assert msg.content == "" + + +def test_tool_message_empty_list_content_normalizes_to_empty_string(): + msg = ChatMessage(role = "tool", content = [], tool_call_id = "call_1") + assert msg.content == "" + + +def test_tool_message_real_content_is_preserved(): + msg = ChatMessage(role = "tool", content = "ok", tool_call_id = "call_1") + assert msg.content == "ok" + + +def test_user_message_still_requires_content(): + with pytest.raises(ValueError): + ChatMessage(role = "user", content = None) + + +def test_assistant_empty_content_still_collapses_to_none(): + msg = ChatMessage(role = "assistant", content = "") + assert msg.content is None diff --git a/studio/backend/tests/test_tool_policy_gates.py b/studio/backend/tests/test_tool_policy_gates.py index 01f6bbbc3f..fad121a4a1 100644 --- a/studio/backend/tests/test_tool_policy_gates.py +++ b/studio/backend/tests/test_tool_policy_gates.py @@ -2,8 +2,8 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. """ -Tests for `_effective_enable_tools` -- the helper that folds the -process-level `tool_policy` over a request's `enable_tools` field. +Tests for `_effective_enable_tools` -- folds the process-level `tool_policy` +over a request's `enable_tools` field. Truth table (policy x payload.enable_tools -> effective): policy=None + payload=None -> None diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index 8b90a46d5a..2ba3310fbe 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -1,9 +1,9 @@ # 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 `_TOOL_XML_RE` (routes/inference.py) -- strips tool-call -XML that leaks past the speculative buffer in core/inference/llama_cpp.py -when the open/close pair is split across the visible/DRAIN boundary. +"""Tests for `_TOOL_XML_RE` (routes/inference.py) -- strips tool-call XML that +leaks past the speculative buffer in core/inference/llama_cpp.py when the +open/close pair is split across the visible/DRAIN boundary. """ from __future__ import annotations @@ -27,11 +27,25 @@ assert _m, "could not extract _TOOL_XML_RE source" _ns = {"_re": _re} exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) _TOOL_XML_RE = _ns["_TOOL_XML_RE"] +_helper = _re.search( + r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n" + r"(?: .+\n)+", + _src, +) +assert _helper, "could not extract _strip_tool_xml_for_display source" +exec(_helper.group(0), _ns) +_strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] # ── Well-formed pairs ───────────────────────────────────────────── +def test_route_display_strip_respects_disabled_auto_heal_contract(): + text = 'literal {"name":"web_search"} survives' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + + def test_strips_well_formed_tool_call(): text = ( "Let me search.\n" @@ -77,9 +91,7 @@ def test_strips_orphan_tool_call_no_close(): def test_strips_orphan_function_no_close(): - text = ( - "I'll call python:\n\n\nprint(1)\n" - ) + text = "I'll call python:\n\n\nprint(1)\n" cleaned = _TOOL_XML_RE.sub("", text) assert "` is required so doc/example prose survives. + # Tail-anchor on `` so doc/example prose survives. text = ( "Here is the Qwen tool-call format:\n" "```xml\n" @@ -209,8 +221,8 @@ def test_real_world_sweep_leaks_get_stripped(leak): # ── Real-world tail-only from gdpval sweep ────────── -# All end-anchored: outer truncated by EOS, -# inner open DRAINED, leaving bare tail. +# All end-anchored: outer truncated by EOS, inner +# open DRAINED, leaving bare tail. GDPVAL_PARAMETER_LEAKS = [ # Qwen3.5-27B Q8_0 / worldbank s00 "the page contains image data and the text is not readable.\n\n\n", diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py index 8ba97af701..1958c8d570 100644 --- a/studio/backend/tests/test_trained_model_scan.py +++ b/studio/backend/tests/test_trained_model_scan.py @@ -28,9 +28,7 @@ from utils.models.model_config import ( ) -def test_scan_trained_models_includes_lora_and_full_finetune_outputs( - tmp_path: Path, monkeypatch -): +def test_scan_trained_models_includes_lora_and_full_finetune_outputs(tmp_path: Path, monkeypatch): # resolve_output_dir refuses absolutes outside outputs_root; point it at tmp_path. from utils.models import model_config as _mc from utils.paths import storage_roots as _sr @@ -53,17 +51,14 @@ def test_scan_trained_models_includes_lora_and_full_finetune_outputs( (full_dir / "model.safetensors").write_bytes(b"") found = { - name: (path, model_type) - for name, path, model_type in scan_trained_models(str(tmp_path)) + name: (path, model_type) for name, path, model_type in scan_trained_models(str(tmp_path)) } assert found[lora_dir.name] == (str(lora_dir), "lora") assert found[full_dir.name] == (str(full_dir), "merged") -def test_get_base_model_from_checkpoint_falls_back_to_full_finetune_config( - tmp_path: Path, -): +def test_get_base_model_from_checkpoint_falls_back_to_full_finetune_config(tmp_path: Path): (tmp_path / "config.json").write_text( json.dumps({"_name_or_path": "HuggingFaceTB/SmolLM-135M"}) ) @@ -85,14 +80,9 @@ def test_get_base_model_from_lora_rejects_full_finetune_dirs(tmp_path: Path): @patch("utils.models.model_config.detect_audio_type", return_value = None) @patch("utils.models.model_config.is_vision_model", return_value = False) def test_model_config_full_finetune_local_path_is_not_lora( - _mock_vision, - _mock_audio_type, - _mock_audio_input, - tmp_path: Path, + _mock_vision, _mock_audio_type, _mock_audio_input, tmp_path: Path ): - (tmp_path / "config.json").write_text( - json.dumps({"_name_or_path": "unsloth/Qwen3-4B"}) - ) + (tmp_path / "config.json").write_text(json.dumps({"_name_or_path": "unsloth/Qwen3-4B"})) (tmp_path / "model.safetensors").write_bytes(b"") config = ModelConfig.from_identifier(str(tmp_path)) diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py new file mode 100644 index 0000000000..a2dc78bee2 --- /dev/null +++ b/studio/backend/tests/test_training_nan_loss_handling.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Pin Studio's behavior when a training event reports non-finite (NaN/Inf) loss. + +The training event handler used to filter NaN/Inf to None silently while +leaving the previous finite loss in progress.loss — so the API kept reporting +the stale value as if everything were fine. We now drop the stale value: +clients see loss=None at the affected step and a one-shot warning is logged. +Training continues; the run is not marked failed. +""" + +from __future__ import annotations + +import math +import os +import sys + +import pytest + +_BACKEND = os.path.join(os.path.dirname(__file__), "..") +if _BACKEND not in sys.path: + sys.path.insert(0, _BACKEND) + +from core.training.training import TrainingBackend + + +def _make_backend() -> TrainingBackend: + return TrainingBackend() + + +def _progress_event( + step: int, + loss: float, + lr: float = 1e-4, +) -> dict: + return { + "type": "progress", + "step": step, + "loss": loss, + "learning_rate": lr, + "epoch": 0.0, + "total_steps": 100, + } + + +class TestNonfiniteLossSoftHandling: + def test_finite_loss_updates_progress_normally(self): + b = _make_backend() + b._handle_event(_progress_event(step = 1, loss = 0.97)) + assert b._progress.loss == pytest.approx(0.97) + assert b._progress.error is None + assert b._should_stop is False + assert getattr(b._progress, "_nonfinite_loss_warned", False) is False + + def test_nan_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step = 1, loss = 0.97)) + assert b._progress.loss == pytest.approx(0.97) + b._handle_event(_progress_event(step = 2, loss = float("nan"))) + # Stale finite loss must NOT leak through + assert b._progress.loss is None + # Run is not marked failed + assert b._progress.error is None + assert b._should_stop is False + # Warning flag is set so we don't re-log on every subsequent NaN step + assert b._progress._nonfinite_loss_warned is True + + def test_inf_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step = 1, loss = float("inf"))) + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + assert b._progress._nonfinite_loss_warned is True + + def test_negative_inf_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step = 1, loss = float("-inf"))) + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + assert b._progress._nonfinite_loss_warned is True + + def test_repeated_nan_only_warns_once(self): + """Subsequent NaN events must not re-fire the warning flag setter. + The flag should already be True after the first NaN.""" + b = _make_backend() + b._handle_event(_progress_event(step = 1, loss = 0.97)) + b._handle_event(_progress_event(step = 2, loss = float("nan"))) + assert b._progress._nonfinite_loss_warned is True + # Further NaN steps don't change anything we care about + b._handle_event(_progress_event(step = 3, loss = float("nan"))) + b._handle_event(_progress_event(step = 4, loss = float("nan"))) + assert b._progress._nonfinite_loss_warned is True + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + + def test_recovery_updates_loss_when_finite_again(self): + """If a NaN step is followed by a finite step, progress.loss must + reflect the new finite value (not stay stuck at None).""" + b = _make_backend() + b._handle_event(_progress_event(step = 1, loss = 0.97)) + b._handle_event(_progress_event(step = 2, loss = float("nan"))) + assert b._progress.loss is None + b._handle_event(_progress_event(step = 3, loss = 0.85)) + assert b._progress.loss == pytest.approx(0.85) + # Warning flag stays set (we don't reset it on recovery) + assert b._progress._nonfinite_loss_warned is True diff --git a/studio/backend/tests/test_training_progress_stream_nan.py b/studio/backend/tests/test_training_progress_stream_nan.py new file mode 100644 index 0000000000..899527a04d --- /dev/null +++ b/studio/backend/tests/test_training_progress_stream_nan.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The SSE progress stream must follow the live progress step during +non-finite-loss stretches (loss reported as null) instead of replaying the +last finite step/loss pair from the metric histories, which skip NaN steps.""" + +import asyncio +import json +import sys +import types + +import pytest + +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.training as rt + + +class _Progress: + def __init__(self): + self.step = 5 + self.total_steps = 10 + self.loss = None # cleared by the NaN honesty fix in core training + self.learning_rate = 8e-5 + self.epoch = 0.1 + self.grad_norm = None + self.num_tokens = None + self.eval_loss = None + self.elapsed_seconds = None + self.eta_seconds = None + + +class _FakeBackend: + """Finite history stops at step 2; live progress is at step 5 with NaN + (loss=None). Active for a few polls, then done.""" + + def __init__(self, active_polls = 2): + self.current_job_id = "job-1" + self.step_history = [1, 2] + self.loss_history = [2.0, 1.5] + self.lr_history = [1e-4, 9e-5] + self.eval_enabled = False + self._active_calls = 0 + self._active_polls = active_polls + self.trainer = types.SimpleNamespace(training_progress = _Progress()) + + def is_training_active(self): + self._active_calls += 1 + return self._active_calls <= self._active_polls + + +class _FakeRequest: + headers = {} + + +def _collect_events(response, timeout = 15): + async def _drain(): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + return "".join(c.decode() if isinstance(c, bytes) else c for c in chunks) + + return asyncio.run(asyncio.wait_for(_drain(), timeout)) + + +def _progress_payloads(raw): + payloads = [] + for block in raw.split("\n\n"): + lines = block.strip().splitlines() + data = next((l[6:] for l in lines if l.startswith("data: ")), None) + if data: + payloads.append(json.loads(data)) + return payloads + + +def test_stream_reports_live_step_with_null_loss_during_nan(monkeypatch): + backend = _FakeBackend(active_polls = 2) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + raw = _collect_events(response) + payloads = _progress_payloads(raw) + assert payloads, f"no SSE payloads parsed from: {raw!r}" + + live = [p for p in payloads if p.get("step") == 5] + assert live, ( + "stream never advanced to the live progress step during the NaN " + f"stretch; steps seen: {[p.get('step') for p in payloads]}" + ) + assert live[0]["loss"] is None + # The stale finite pair must not be re-emitted as the latest progress. + stale = [p for p in payloads if p.get("step") == 2 and p.get("loss") == 1.5] + assert not stale + + +def test_inactive_stream_completes_with_live_step_and_null_loss(monkeypatch): + # Fresh connection after the run already ended during a NaN stretch: the + # immediate complete event must not replay the stale finite pair either. + backend = _FakeBackend(active_polls = 0) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + payloads = _progress_payloads(_collect_events(response)) + final = payloads[-1] + assert final["step"] == 5 + assert final["loss"] is None + + +def test_stream_uses_finite_history_when_progress_in_sync(monkeypatch): + backend = _FakeBackend(active_polls = 2) + # Live progress agrees with the history tail: normal finite behavior. + backend.trainer.training_progress.step = 2 + backend.trainer.training_progress.loss = 1.5 + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + payloads = _progress_payloads(_collect_events(response)) + finite = [p for p in payloads if p.get("step") == 2] + assert finite and finite[0]["loss"] == 1.5 diff --git a/studio/backend/tests/test_training_raw_support.py b/studio/backend/tests/test_training_raw_support.py index 384247a191..2b1299de5c 100644 --- a/studio/backend/tests/test_training_raw_support.py +++ b/studio/backend/tests/test_training_raw_support.py @@ -214,10 +214,7 @@ class TestTrainingRawSupport(unittest.TestCase): self.assertEqual(result.dataset[0]["text"], "hello") self.assertEqual(result.dataset[1]["text"], "world") self.assertTrue( - any( - "null or non-string 'text' values" in notice.message - for notice in result.notices - ) + any("null or non-string 'text' values" in notice.message for notice in result.notices) ) diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 94279c28b4..3c5d6cd094 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -15,7 +15,13 @@ from core.training import worker def _missing_flash_attn_import(): real_import = builtins.__import__ - def fake_import(name, globals = None, locals = None, fromlist = (), level = 0): + def fake_import( + name, + globals = None, + locals = None, + fromlist = (), + level = 0, + ): if name == "flash_attn": raise ImportError return real_import(name, globals, locals, fromlist, level) @@ -26,7 +32,13 @@ def _missing_flash_attn_import(): def _missing_module_import(missing: str): real_import = builtins.__import__ - def fake_import(name, globals = None, locals = None, fromlist = (), level = 0): + def fake_import( + name, + globals = None, + locals = None, + fromlist = (), + level = 0, + ): if name == missing: raise ImportError return real_import(name, globals, locals, fromlist, level) @@ -37,9 +49,7 @@ def _missing_module_import(missing: str): def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch): monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) assert worker._should_try_runtime_flash_attn_install(32767) is False - assert worker._should_try_runtime_flash_attn_install( - 32768 - ) is sys.platform.startswith("linux") + assert worker._should_try_runtime_flash_attn_install(32768) is sys.platform.startswith("linux") monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1") assert worker._should_try_runtime_flash_attn_install(32768) is False @@ -105,7 +115,12 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): ) monkeypatch.setattr(worker, "install_wheel", mock.Mock()) - def fake_run(cmd, stdout = None, stderr = None, text = None): + def fake_run( + cmd, + stdout = None, + stderr = None, + text = None, + ): calls.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0, "") @@ -131,9 +146,7 @@ def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): 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, "_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( @@ -208,7 +221,7 @@ def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch): def _force_missing_fla_imports(monkeypatch): - """Make fla.modules / fla.ops.gated_delta_rule imports raise ImportError.""" + """Force fla.modules / fla.ops imports to raise ImportError.""" real_import = builtins.__import__ def fake_import(name, *a, **kw): @@ -254,8 +267,8 @@ def test_flash_linear_attention_skips_for_unrelated_models(monkeypatch): def test_flash_linear_attention_skips_for_ssm_only_models(monkeypatch): - # Nemotron-H / Falcon-H1 / Granite-H / LFM2 take the mamba_ssm path - # and never call FLA's gated_delta_rule kernels. + # Nemotron-H / Falcon-H1 / Granite-H / LFM2 take the mamba_ssm path, + # never FLA's gated_delta_rule kernels. run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) monkeypatch.setattr(worker._sp, "run", run_mock) @@ -276,7 +289,7 @@ def test_flash_linear_attention_matches_full_qwen3_family(monkeypatch): monkeypatch.setattr(worker._sp, "run", run_mock) _force_missing_fla_imports(monkeypatch) monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) - # Hermetic discovery: pretend installed transformers ships all the Qwen GDN families. + # Hermetic discovery: pretend transformers ships all Qwen GDN families. monkeypatch.setattr( worker, "_discover_fla_model_types", @@ -357,9 +370,8 @@ def test_flash_linear_attention_install_includes_einops(monkeypatch): args = run_mock.call_args[0][0] assert "--no-deps" in args - # einops is declared by fla-core; packaging and triton are pulled in - # because fla/utils.py imports them at module load but neither is - # declared in fla-core's METADATA (an upstream FLA gap). + # packaging and triton are added because fla/utils.py imports them at load + # but neither is in fla-core's METADATA (an upstream FLA gap). assert "einops" in args assert "packaging" in args assert "triton" in args @@ -376,8 +388,8 @@ def test_flash_linear_attention_logs_post_install_import_failure(monkeypatch): def fake_importable(): import_calls["count"] += 1 - # First call (pre-install probe) -> False so we attempt install. - # Second call (post-install verify) -> still False. + # Pre-install probe -> False (attempt install); post-install + # verify -> still False. return False monkeypatch.setattr(worker, "_flash_linear_attention_importable", fake_importable) @@ -420,13 +432,13 @@ def test_tilelang_backend_pins_only_binary(monkeypatch): run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) monkeypatch.setattr(worker._sp, "run", run_mock) monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) - # Need to bypass the post-install probe too. + # Bypass the post-install probe too. probe_calls = {"count": 0} def fake_probe(): probe_calls["count"] += 1 - # First probe (pre-install): False so install runs. - # Second probe (post-install): True so success branch taken. + # Pre-install probe: False (install runs); post-install: True + # (success branch taken). return probe_calls["count"] > 1 monkeypatch.setattr(worker, "_tilelang_importable", fake_probe) @@ -478,14 +490,10 @@ def test_tilelang_backend_installs_pinned_pair_for_qwen3_5(monkeypatch): def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch): """Repair path issues TWO pip calls: - Call 1 (repair): `--force-reinstall --no-deps apache-tvm-ffi==0.1.9` - — surgically downgrades the broken package only. `--no-deps` here - is REQUIRED to prevent --force-reinstall from cascading through - apache-tvm-ffi's dep graph and replacing torch / the CUDA stack. - - Call 2 (install): plain `apache-tvm-ffi==0.1.9 tilelang==0.1.8` - — resolves missing transitive deps (z3-solver, ml-dtypes) without - --force-reinstall, so it never replaces already-correct packages. + 1 (repair): --force-reinstall --no-deps apache-tvm-ffi -- downgrades only + the broken package; --no-deps stops the cascade through its deps to torch. + 2 (install): plain apache-tvm-ffi + tilelang -- resolves missing transitive + deps without --force-reinstall, so it never replaces correct packages. """ monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") @@ -502,18 +510,14 @@ def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch): assert run_mock.call_count == 2 repair_args, install_args = (call[0][0] for call in run_mock.call_args_list) - # Repair: --force-reinstall --no-deps, apache-tvm-ffi ONLY (no tilelang). + # Repair: --force-reinstall --no-deps, apache-tvm-ffi ONLY. assert "--force-reinstall" in repair_args - assert ( - "--no-deps" in repair_args - ), "Repair MUST use --no-deps to avoid replacing torch / CUDA" + assert "--no-deps" in repair_args, "Repair MUST use --no-deps to avoid replacing torch / CUDA" assert "--only-binary=:all:" in repair_args assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in repair_args - assert all( - "tilelang" not in a for a in repair_args - ), "Repair MUST only touch apache-tvm-ffi" + assert all("tilelang" not in a for a in repair_args), "Repair MUST only touch apache-tvm-ffi" - # Install: regular dep-resolving install, NO --force-reinstall. + # Install: regular dep-resolving install, no --force-reinstall. assert "--force-reinstall" not in install_args assert "--no-deps" not in install_args assert "--only-binary=:all:" in install_args @@ -564,7 +568,7 @@ def test_tilelang_backend_swallows_install_timeout(monkeypatch): statuses: list[str] = [] monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg)) - # Should not raise. + # Must not raise. worker._ensure_tilelang_backend( event_queue = [], model_name = "unsloth/Qwen3.5-2B", @@ -579,7 +583,7 @@ def test_tilelang_backend_skipped_for_ssm_models(monkeypatch): monkeypatch.setattr(worker._sp, "run", run_mock) # Nemotron-H / Falcon-H1 / Granite-H take the mamba_ssm path, not FLA's - # gated_delta_rule -> tilelang has no effect on them. + # gated_delta_rule -> tilelang doesn't affect them. for name in ( "tiiuae/Falcon-H1-0.5B-Instruct", "nvidia/Nemotron-H-8B-Base", @@ -624,27 +628,23 @@ def test_tilelang_backend_swallows_install_failure(monkeypatch): assert any("failed" in s.lower() for s in statuses) -# ─────────────────────────────────────────────────────────────────── -# Runtime hook on `is_flash_linear_attention_available` / -# `is_causal_conv1d_available`. These are the primary gate in -# normal operation; the substring tests above cover the -# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 fallback. -# ─────────────────────────────────────────────────────────────────── +# Runtime hook on is_flash_linear_attention_available / +# is_causal_conv1d_available -- the primary gate in normal operation. The +# substring tests above cover the SKIP_FAST_PATH_HOOKS=1 fallback. class _FakeQueue(list): - """List with `.put` so worker._send_status can send into it during tests.""" + """List with `.put` so worker._send_status can send into it in tests.""" def put(self, item): self.append(item) def _make_fake_gate(initial_return: bool): - """Build a callable that mimics transformers' lru_cache-decorated gates. + """Callable mimicking transformers' lru_cache-decorated gates. - Tracks call count and exposes a `cache_clear` attribute. The return - value can be flipped to mimic install-then-True behaviour by setting - `.next_return`. + Tracks call count and exposes `cache_clear`. Flip `.next_return` to + mimic install-then-True behaviour. """ class Gate: @@ -689,20 +689,16 @@ def test_hook_installs_when_gate_returns_false(monkeypatch): conv_install = mock.Mock(side_effect = _conv_install_side_effect) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fla_install - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu - # Both gates are now wrapped. Call them — the hook should drive the install. + # Both gates wrapped; calling them should drive the install. assert _iu.is_flash_linear_attention_available() is True fla_install.assert_called_once() tile_install.assert_called_once() @@ -711,9 +707,9 @@ def test_hook_installs_when_gate_returns_false(monkeypatch): def test_hook_skips_install_when_gate_already_true(monkeypatch): - """When both gates are already True AND tilelang is healthy, the hook - must do zero install work. (Tilelang repair on the already-True path - is covered by test_hook_runs_tilelang_repair_when_fla_already_true.) + """Both gates already True AND tilelang healthy -> zero install work. + (Tilelang repair on the already-True path is covered by + test_hook_runs_tilelang_repair_when_fla_already_true.) """ fla_gate = _make_fake_gate(initial_return = True) conv_gate = _make_fake_gate(initial_return = True) @@ -722,21 +718,16 @@ def test_hook_skips_install_when_gate_already_true(monkeypatch): fla_install = mock.Mock() tile_install = mock.Mock() conv_install = mock.Mock() - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fla_install - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) - # Tilelang healthy so the post_available path is a no-op (otherwise - # it would call tile_install, which is correct behaviour but - # outside the scope of this test). + # Tilelang healthy -> post_available path is a no-op (otherwise it + # would call tile_install, correct but out of scope here). monkeypatch.setattr(worker, "_tilelang_importable", lambda: True) monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.9") monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu @@ -764,22 +755,18 @@ def test_hook_idempotent_on_repeat_call(monkeypatch): return True conv_install = mock.Mock(side_effect = _conv_install_side_effect) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fla_install - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu # First call: hook fires. _iu.is_flash_linear_attention_available() - # Subsequent calls: must not re-trigger the installer. + # Later calls: must not re-trigger the installer. _iu.is_flash_linear_attention_available() _iu.is_flash_linear_attention_available() assert fla_install.call_count == 1 @@ -794,22 +781,16 @@ def test_hook_handles_install_failure_gracefully(monkeypatch): def raising_install(eq): raise RuntimeError("pip failed to fetch wheel") - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", raising_install - ) - monkeypatch.setattr( - worker, "_ensure_tilelang_backend_unconditional", lambda eq: None - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", raising_install) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: None) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: None) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu - # Must not raise; returns False so transformers falls back to torch loop. + # Must not raise; returns False so transformers uses the torch loop. assert _iu.is_flash_linear_attention_available() is False @@ -819,18 +800,14 @@ def test_hook_can_be_disabled_via_env(monkeypatch): _patch_iu_gates(monkeypatch, fla_gate, conv_gate) fla_install = mock.Mock() - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fla_install - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) monkeypatch.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1") - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu - # Hook should NOT have been installed; gates remain the fakes. + # Hook not installed; gates remain the fakes. assert _iu.is_flash_linear_attention_available is fla_gate assert _iu.is_causal_conv1d_available is conv_gate fla_install.assert_not_called() @@ -841,36 +818,29 @@ def test_hook_clears_lru_cache_before_first_check(monkeypatch): conv_gate = _make_fake_gate(initial_return = True) _patch_iu_gates(monkeypatch, fla_gate, conv_gate) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", lambda eq: None - ) - monkeypatch.setattr( - worker, "_ensure_tilelang_backend_unconditional", lambda eq: None - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: None) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: None) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: None) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu _iu.is_flash_linear_attention_available() - # The wrapper called cache_clear at least once before delegating. + # Wrapper called cache_clear at least once before delegating. assert fla_gate.cache_clear_count >= 1 def test_hook_rewrites_previously_imported_module_bindings(monkeypatch): - """Modeling files bind `is_flash_linear_attention_available` locally - via `from ... import is_X`. Reassigning the attribute on - transformers.utils.import_utils alone does NOT reach those local - bindings. The hook installer sweeps sys.modules and rebinds them. + """Modeling files bind is_flash_linear_attention_available locally via + `from ... import is_X`. Reassigning the attribute on import_utils alone + misses those; the hook installer sweeps sys.modules and rebinds them. """ fla_gate = _make_fake_gate(initial_return = False) conv_gate = _make_fake_gate(initial_return = True) _patch_iu_gates(monkeypatch, fla_gate, conv_gate) - # Create a fake modeling module that did `from ... import is_flash_linear_attention_available`. + # Fake modeling module that did `from ... import is_flash_linear_attention_available`. fake_mod = sys.modules.setdefault( "_test_fake_modeling_qwen35", type(sys)("_test_fake_modeling_qwen35") ) @@ -880,22 +850,16 @@ def test_hook_rewrites_previously_imported_module_bindings(monkeypatch): fla_gate.next_return = True return True - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fake_install - ) - monkeypatch.setattr( - worker, "_ensure_tilelang_backend_unconditional", lambda eq: True - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fake_install) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") - # The fake module's local binding has been rewritten to the wrapper. + # The fake module's local binding is rewritten to the wrapper. assert fake_mod.is_flash_linear_attention_available is not fla_gate - # Calling through the fake module's reference triggers the install. + # Calling through the fake module's reference triggers install. assert fake_mod.is_flash_linear_attention_available() is True del sys.modules["_test_fake_modeling_qwen35"] @@ -915,35 +879,24 @@ def test_hook_skips_when_import_utils_unavailable(monkeypatch): monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) # Should not raise. - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch): - """Hook disabled -> legacy gate falls back to auto-discovered model types.""" + """Hook disabled -> legacy gate falls back to auto-discovered types.""" install_mock = mock.Mock() - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", install_mock - ) - monkeypatch.setattr( - worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_5"}) - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", install_mock) + monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_5"})) monkeypatch.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1") - worker._ensure_flash_linear_attention( - event_queue = [], model_name = "unsloth/Qwen3.5-2B" - ) + worker._ensure_flash_linear_attention(event_queue = [], model_name = "unsloth/Qwen3.5-2B") assert install_mock.call_count == 1 - worker._ensure_flash_linear_attention( - event_queue = [], model_name = "meta-llama/Llama-3.1-8B" - ) + worker._ensure_flash_linear_attention(event_queue = [], model_name = "meta-llama/Llama-3.1-8B") assert install_mock.call_count == 1 -# ─────────────────────────────────────────────────────────────────── -# Regression tests for the 10-reviewer findings: +# Regression tests for the reviewer findings: # 1. tilelang Qwen-guard on hook path (non-Qwen FLA models) # 2. tilelang repair must not replace torch / CUDA stack # 3. hook must trust installer's bool, not transformers metadata @@ -952,12 +905,11 @@ def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch): # 6. tilelang skipped when FLA was skipped / failed # 7. tilelang repair runs when FLA is already True # 8. older FLA detected as stale and reinstalled -# ─────────────────────────────────────────────────────────────────── def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch): - """A model whose name is not in the auto-discovered FLA allowlist calls - is_flash_linear_attention_available but should NOT get tilelang.""" + """A model not in the auto-discovered FLA allowlist calls + is_flash_linear_attention_available but must NOT get tilelang.""" fla_gate = _make_fake_gate(initial_return = False) conv_gate = _make_fake_gate(initial_return = True) _patch_iu_gates(monkeypatch, fla_gate, conv_gate) @@ -968,17 +920,13 @@ def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch) fla_install = mock.Mock(side_effect = _fla_install) tile_install = mock.Mock(return_value = True) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fla_install - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) - monkeypatch.setattr( - worker, "_install_package_wheel_first", mock.Mock(return_value = True) - ) + monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) # Hermetize the auto-discovered set so the test stays valid as new # transformers releases add FLA-using model_types (eg olmo_hybrid in - # 5.4.0). The semantic under test is "outside-allowlist -> no tilelang". + # 5.4.0). Test semantic: "outside-allowlist -> no tilelang". monkeypatch.setattr( worker, "_discover_fla_model_types", @@ -1009,18 +957,12 @@ def test_hook_does_install_tilelang_for_qwen35(monkeypatch): fla_install = mock.Mock(side_effect = _fla_install) tile_install = mock.Mock(return_value = True) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fla_install - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) - monkeypatch.setattr( - worker, "_install_package_wheel_first", mock.Mock(return_value = True) - ) + monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu @@ -1031,7 +973,7 @@ def test_hook_does_install_tilelang_for_qwen35(monkeypatch): def test_tilelang_repair_does_not_touch_torch_cuda_stack(monkeypatch): """Finding #2: the broken-tvm-ffi repair must use --no-deps on the - forced step so --force-reinstall does not cascade through + forced step so --force-reinstall doesn't cascade through apache-tvm-ffi's dep graph and pull a different torch wheel. """ monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) @@ -1045,54 +987,38 @@ def test_tilelang_repair_does_not_touch_torch_cuda_stack(monkeypatch): assert run_mock.call_count == 2 repair_args = run_mock.call_args_list[0][0][0] - # The forced step MUST be --no-deps so torch / CUDA stack is untouched. + # Forced step MUST be --no-deps so torch / CUDA stack is untouched. assert "--force-reinstall" in repair_args and "--no-deps" in repair_args - # And it touches ONLY apache-tvm-ffi, not tilelang / torch. + # Touches ONLY apache-tvm-ffi, not tilelang / torch. assert all("tilelang" not in a for a in repair_args) assert all("torch" not in a for a in repair_args) def test_hook_trusts_installer_bool_not_metadata(monkeypatch): - """Finding #3: if pip exits 0 but deep imports fail, the installer - returns False; the hook must propagate False even if the underlying - `original()` gate (which only checks metadata) returns True after - pip succeeds. - - Setup mirrors the real bug: - 1. Pre-install: gate=False (FLA not present) → wrapper triggers install. - 2. Installer's `_flash_linear_attention_importable` post-probe fails, - so the installer returns False. (pip exited 0 but `import fla.modules` - raised because of a missing transitive dep.) - 3. Post-install: gate would return True (metadata check sees fla-core - version) — but the wrapper must IGNORE that and use the installer's - False so transformers takes the torch fallback. + """Finding #3: if pip exits 0 but deep imports fail, the installer returns + False; the hook must propagate that False even if the metadata-only gate + returns True after pip succeeds, so transformers takes the torch fallback. """ # Gate flips True after install (simulating "metadata sees fla"). fla_gate = _make_fake_gate(initial_return = False) conv_gate = _make_fake_gate(initial_return = True) _patch_iu_gates(monkeypatch, fla_gate, conv_gate) - # Installer "succeeds" at pip, AND flips the gate to True (metadata - # sees fla post-install), BUT returns False (deep import broken). + # Installer "succeeds" at pip and flips the gate to True (metadata + # sees fla post-install), but returns False (deep import broken). def _bad_install(eq): fla_gate.next_return = True # metadata says yes after pip return False # but deep import is broken fake_fla_install = mock.Mock(side_effect = _bad_install) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fake_fla_install - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fake_fla_install) monkeypatch.setattr( worker, "_ensure_tilelang_backend_unconditional", mock.Mock(return_value = True) ) - monkeypatch.setattr( - worker, "_install_package_wheel_first", mock.Mock(return_value = True) - ) + monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu @@ -1102,9 +1028,9 @@ def test_hook_trusts_installer_bool_not_metadata(monkeypatch): def test_rebind_does_not_trigger_module_getattr(monkeypatch): - """Finding #5: the rebind sweep must use __dict__, not getattr(), - to avoid invoking transformers' lazy module __getattr__ which spits - out hundreds of "Accessing X from .models..." warnings. + """Finding #5: the rebind sweep must use __dict__, not getattr(), to + avoid invoking transformers' lazy module __getattr__ which spits out + hundreds of "Accessing X from .models..." warnings. """ original = object() replacement = object() @@ -1119,8 +1045,8 @@ def test_rebind_does_not_trigger_module_getattr(monkeypatch): lazy = _GetattrTripwire("_lazy_test_module") sys.modules["_lazy_test_module"] = lazy try: - # No module-level binding to `is_flash_linear_attention_available` - # in __dict__, so the sweep must NOT trip the tripwire. + # No `is_flash_linear_attention_available` in __dict__, so the + # sweep must NOT trip the tripwire. worker._rebind_in_already_imported_modules( attr_name = "is_flash_linear_attention_available", old_obj = original, @@ -1136,7 +1062,7 @@ def test_rebind_does_not_trigger_module_getattr(monkeypatch): def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch): """Finding #6: env-skipped FLA returns False from _ensure_flash_linear_attention_unconditional; tilelang must NOT - install in that case. + install then. """ fla_gate = _make_fake_gate(initial_return = False) conv_gate = _make_fake_gate(initial_return = True) @@ -1145,14 +1071,10 @@ def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch): monkeypatch.setenv(worker._FLA_SKIP_ENV, "1") tile_install = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) - monkeypatch.setattr( - worker, "_install_package_wheel_first", mock.Mock(return_value = True) - ) + monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu @@ -1162,9 +1084,9 @@ def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch): def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch): - """Finding #7: when FLA is already importable (gate returns True at - first probe) but tilelang is missing or apache-tvm-ffi is on the - broken list, the post-available action must still run tilelang. + """Finding #7: when FLA is already importable (gate True at first + probe) but tilelang is missing or apache-tvm-ffi is on the broken + list, the post-available action must still run tilelang. """ fla_gate = _make_fake_gate(initial_return = True) conv_gate = _make_fake_gate(initial_return = True) @@ -1172,38 +1094,32 @@ def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch): fla_install = mock.Mock(return_value = True) tile_install = mock.Mock(return_value = True) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", fla_install - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) - monkeypatch.setattr( - worker, "_install_package_wheel_first", mock.Mock(return_value = True) - ) - # tilelang missing AND tvm-ffi is on broken list — both trigger repair. + monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) + # tilelang missing AND tvm-ffi on broken list — both trigger repair. monkeypatch.setattr(worker, "_tilelang_importable", lambda: False) monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.11") monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") from transformers.utils import import_utils as _iu _iu.is_flash_linear_attention_available() - # FLA install was NOT needed; tilelang repair WAS still triggered. + # FLA install NOT needed; tilelang repair still triggered. fla_install.assert_not_called() tile_install.assert_called_once() def test_fla_installer_force_reinstalls_when_older_version_present(monkeypatch): - """Finding #8: when an older `flash-linear-attention` is importable - but below the pin, the installer must force a reinstall (not no-op). + """Finding #8: an older `flash-linear-attention` that is importable + but below the pin must force a reinstall (not no-op). """ monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False) monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9)) - # Importable but stale (current() reports False even though importable() is True). + # Importable but stale (current()=False though importable()=True). monkeypatch.setattr(worker, "_flash_linear_attention_importable", lambda: True) monkeypatch.setattr(worker, "_flash_linear_attention_current", lambda **kw: False) run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) @@ -1222,23 +1138,19 @@ def test_fla_installer_force_reinstalls_when_older_version_present(monkeypatch): def test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode(): - """Finding #4: SSM modeling files use `lazy_load_kernel("causal-conv1d")` - and never call `is_causal_conv1d_available()`, so the hook would not - fire for them. The orchestrator must always run the eager - substring installer regardless of hook mode. - - This test reads the worker source rather than running the full - orchestrator (which requires a configured training config). It - asserts the eager install is OUTSIDE the if/else hook branch. + """Finding #4: SSM modeling files use lazy_load_kernel and never call + is_causal_conv1d_available(), so the hook won't fire; the orchestrator must + always run the eager installer regardless of hook mode. Reads the worker + source and asserts the eager install is OUTSIDE the if/else hook branch. """ import inspect src = inspect.getsource(worker.run_training_process) - # Find the orchestration block. + # Orchestration block. assert "_ensure_causal_conv1d_fast_path(event_queue, model_name)" in src assert "_install_fast_path_hooks(event_queue, model_name)" in src - # The eager causal_conv1d call must appear BEFORE the hook-mode if/else, - # not nested inside the `if _FAST_PATH_HOOKS_SKIP_ENV` branch. + # Eager causal_conv1d call must come BEFORE the hook-mode if/else, not + # nested inside the `if _FAST_PATH_HOOKS_SKIP_ENV` branch. eager_pos = src.find("_ensure_causal_conv1d_fast_path(event_queue, model_name)") skip_check_pos = src.find('os.getenv(_FAST_PATH_HOOKS_SKIP_ENV) == "1"') assert eager_pos < skip_check_pos, ( @@ -1248,19 +1160,16 @@ def test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode(): ) -# ─────────────────────────────────────────────────────────────────── -# HIP / ROCm regression coverage (h34v3nzc0dex Strix Halo report). -# tilelang 0.1.8 has no HIP GEMM backend; FLA's TileLang dispatch -# crashes mid-backward on AMD with "Unsupported target for gemm: hip". -# The fix: skip the install on HIP-built torch AND setdefault -# FLA_TILELANG=0 so already-installed tilelang doesn't get used either. -# ─────────────────────────────────────────────────────────────────── +# HIP / ROCm regression coverage (Strix Halo report). +# tilelang 0.1.8 has no HIP GEMM backend; FLA's TileLang dispatch crashes +# mid-backward on AMD ("Unsupported target for gemm: hip"). Fix: skip install on +# HIP torch AND setdefault FLA_TILELANG=0 so an existing tilelang isn't used. def test_tilelang_platform_unsupported_on_hip_torch(monkeypatch): - """Strix Halo / MI300 with ROCm torch: linux + x86_64 looks - identical to a CUDA box at the OS level, so the platform check - must consult torch.version.hip explicitly. + """Strix Halo / MI300 with ROCm torch: linux + x86_64 looks identical + to a CUDA box at the OS level, so the platform check must consult + torch.version.hip explicitly. """ monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) assert worker._tilelang_platform_supported() is False @@ -1281,50 +1190,38 @@ def test_tilelang_install_skipped_on_hip_torch(monkeypatch): def test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip(monkeypatch): - """When HIP torch is detected, hook installer must set - FLA_TILELANG=0 (via setdefault — respects user override) so any - PRE-EXISTING tilelang install isn't used by FLA's dispatcher. + """On HIP torch, the hook installer must setdefault FLA_TILELANG=0 + (respecting user override) so a PRE-EXISTING tilelang install isn't + used by FLA's dispatcher. """ import os as _os monkeypatch.delenv("FLA_TILELANG", raising = False) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True - ) - monkeypatch.setattr( - worker, "_ensure_tilelang_backend_unconditional", lambda eq: True - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") assert _os.environ.get("FLA_TILELANG") == "0" def test_install_fast_path_hooks_respects_user_fla_tilelang_override(monkeypatch): - """If the user explicitly set FLA_TILELANG (even on HIP), don't - overwrite — they may know they have a HIP-aware tilelang fork. + """If the user set FLA_TILELANG (even on HIP), don't overwrite — they + may have a HIP-aware tilelang fork. """ import os as _os monkeypatch.setenv("FLA_TILELANG", "1") monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True - ) - monkeypatch.setattr( - worker, "_ensure_tilelang_backend_unconditional", lambda eq: True - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") assert _os.environ["FLA_TILELANG"] == "1" @@ -1336,17 +1233,11 @@ def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch): monkeypatch.delenv("FLA_TILELANG", raising = False) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) monkeypatch.setattr(worker, "_torch_has_hip", lambda: False) - monkeypatch.setattr( - worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True - ) - monkeypatch.setattr( - worker, "_ensure_tilelang_backend_unconditional", lambda eq: True - ) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) - worker._install_fast_path_hooks( - event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" - ) + worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") assert _os.environ.get("FLA_TILELANG") is None @@ -1356,10 +1247,8 @@ def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch): # ─────────────────────────────────────────────────────────────────── -def _make_fake_transformers_tree( - tmp_path, fla_types: list[str], non_fla_types: list[str] -): - """Lay out a tmp dir as `transformers/models/{type}/modeling_{type}.py`.""" +def _make_fake_transformers_tree(tmp_path, fla_types: list[str], non_fla_types: list[str]): + """Lay out tmp dir as `transformers/models/{type}/modeling_{type}.py`.""" pkg = tmp_path / "transformers" models = pkg / "models" models.mkdir(parents = True) @@ -1401,9 +1290,7 @@ def test_discover_fla_model_types_returns_only_fla_users(tmp_path, monkeypatch): def test_discover_fla_model_types_caches_across_calls(tmp_path, monkeypatch): - pkg = _make_fake_transformers_tree( - tmp_path, fla_types = ["qwen3_5"], non_fla_types = [] - ) + pkg = _make_fake_transformers_tree(tmp_path, fla_types = ["qwen3_5"], non_fla_types = []) fake = mock.MagicMock(__file__ = str(pkg / "__init__.py")) monkeypatch.setitem(sys.modules, "transformers", fake) _reset_fla_cache(monkeypatch) @@ -1424,7 +1311,7 @@ def test_discover_fla_model_types_caches_across_calls(tmp_path, monkeypatch): second = worker._discover_fla_model_types() assert first == second - assert read_calls[0] == after_first # cache hit: no extra disk reads + assert read_calls[0] == after_first # cache hit: no extra reads def test_discover_fla_model_types_handles_missing_transformers(monkeypatch): @@ -1432,7 +1319,13 @@ def test_discover_fla_model_types_handles_missing_transformers(monkeypatch): real_import = builtins.__import__ - def fake_import(name, globals = None, locals = None, fromlist = (), level = 0): + def fake_import( + name, + globals = None, + locals = None, + fromlist = (), + level = 0, + ): if name == "transformers": raise ImportError("transformers not installed") return real_import(name, globals, locals, fromlist, level) @@ -1443,9 +1336,7 @@ def test_discover_fla_model_types_handles_missing_transformers(monkeypatch): def test_discover_fla_model_types_handles_unreadable_file(tmp_path, monkeypatch): - pkg = _make_fake_transformers_tree( - tmp_path, fla_types = ["qwen3_5"], non_fla_types = [] - ) + pkg = _make_fake_transformers_tree(tmp_path, fla_types = ["qwen3_5"], non_fla_types = []) fake = mock.MagicMock(__file__ = str(pkg / "__init__.py")) monkeypatch.setitem(sys.modules, "transformers", fake) _reset_fla_cache(monkeypatch) @@ -1461,7 +1352,7 @@ def test_discover_fla_model_types_handles_unreadable_file(tmp_path, monkeypatch) monkeypatch.setattr(_Path, "read_text", boom_read) result = worker._discover_fla_model_types() - assert result == frozenset() # unreadable file simply doesn't contribute + assert result == frozenset() # unreadable file doesn't contribute def test_model_wants_tilelang_handles_real_repo_names(monkeypatch): @@ -1491,9 +1382,7 @@ def test_model_wants_tilelang_empty_when_transformers_has_no_fla(monkeypatch): def test_model_wants_tilelang_normalizes_separators(monkeypatch): - monkeypatch.setattr( - worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_next"}) - ) + monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_next"})) for variant in ( "qwen3-next", "Qwen3.Next", @@ -1504,22 +1393,17 @@ def test_model_wants_tilelang_normalizes_separators(monkeypatch): assert worker._model_wants_tilelang(variant) is True, variant -# ──────────────────────────────────────────────────────────────────── -# HIP source-build gcc-install-dir coverage (h34v3nzc0dex Strix Halo). -# Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14, -# so ROCm clang-20 picks it and fails with 'cstdlib' file not found -# when building causal-conv1d (or any other HIP source fallback). -# _hipcc_gcc_install_dir() finds a gcc dir that has both halves; the -# _install_package_wheel_first HIP branch passes it to clang via -# HIPCC_COMPILE_FLAGS_APPEND. Parallel to bbf004c's setup.sh fix for -# the llama.cpp HIP build (PR #5301). -# ──────────────────────────────────────────────────────────────────── +# HIP source-build gcc-install-dir coverage (Strix Halo). +# Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14, so ROCm +# clang-20 picks it and fails ('cstdlib' not found) building causal-conv1d. +# _hipcc_gcc_install_dir() finds a gcc dir with both halves; the HIP branch of +# _install_package_wheel_first passes it via HIPCC_COMPILE_FLAGS_APPEND. +# Parallels bbf004c's setup.sh fix for the llama.cpp HIP build (PR #5301). def _isdir_for_layout(*existing: str): - """Return an os.path.isdir replacement that only treats the given - absolute paths as directories. Lets a test simulate exactly which - gcc runtime dirs and C++ header dirs exist on the host.""" + """os.path.isdir replacement treating only the given absolute paths as + directories, to simulate which gcc runtime / C++ header dirs exist.""" valid = set(existing) def fake_isdir(path: str) -> bool: @@ -1530,7 +1414,7 @@ def _isdir_for_layout(*existing: str): def test_hipcc_gcc_install_dir_picks_highest_with_headers(monkeypatch): """gcc-14 has runtime but no /usr/include/c++/14; loop falls through - to gcc-13 which has both. This is the exact Ubuntu 24.04 layout.""" + to gcc-13 which has both. The exact Ubuntu 24.04 layout.""" monkeypatch.setattr(sys, "platform", "linux") import platform as _platform @@ -1566,8 +1450,8 @@ def test_hipcc_gcc_install_dir_picks_14_when_headers_exist(monkeypatch): def test_hipcc_gcc_install_dir_returns_none_when_no_match(monkeypatch): - """No gcc dir has both halves → return None and skip the env injection - rather than guessing wrong and surfacing a confusing build failure.""" + """No gcc dir has both halves → return None and skip env injection + rather than guessing wrong and causing a confusing build failure.""" monkeypatch.setattr(sys, "platform", "linux") import platform as _platform @@ -1597,10 +1481,9 @@ def test_hipcc_gcc_install_dir_returns_none_on_non_x86_64(monkeypatch): def _make_hip_install_env(monkeypatch, *, gcc_dir: str | None): - """Common scaffolding for tests that exercise the HIP source-build - branch of _install_package_wheel_first end-to-end. The package isn't - installed yet, no prebuilt wheel exists, hipcc is on PATH, and the - fake env reports an HIP torch.""" + """Scaffolding for end-to-end tests of the HIP source-build branch of + _install_package_wheel_first: package not installed, no prebuilt + wheel, hipcc on PATH, fake env reports HIP torch.""" monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d")) monkeypatch.setattr( worker, @@ -1655,8 +1538,8 @@ def test_install_injects_gcc_install_dir_on_hip_source_build(monkeypatch): def test_install_appends_to_existing_hipcc_compile_flags(monkeypatch): - """User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' set → final value - keeps the user's flags AND adds --gcc-install-dir at the end.""" + """User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' → final value keeps + the user's flags AND appends --gcc-install-dir.""" monkeypatch.setenv("HIPCC_COMPILE_FLAGS_APPEND", "-O3 -DFOO") _make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13") @@ -1717,9 +1600,9 @@ def test_install_respects_user_gcc_install_dir(monkeypatch): release_base_url = "https://example.com", ) - # subprocess.run was invoked without env override (the user already - # set HIPCC_COMPILE_FLAGS_APPEND with --gcc-install-dir, so we left - # the env alone — the existing value is inherited normally). + # subprocess.run invoked without env override (user already set + # HIPCC_COMPILE_FLAGS_APPEND with --gcc-install-dir, so we left the + # env alone — the existing value is inherited). assert captured == {"_called": "yes_no_env"} @@ -1741,7 +1624,7 @@ def test_install_does_not_inject_env_on_cuda(monkeypatch): monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None) monkeypatch.setattr(worker.shutil, "which", lambda name: None) monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) - # If _hipcc_gcc_install_dir were called on CUDA we'd want to know. + # _hipcc_gcc_install_dir must not be called on CUDA. monkeypatch.setattr( worker, "_hipcc_gcc_install_dir", diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index c031c2fea3..7c497ba1b1 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -10,9 +10,8 @@ from unittest.mock import patch # --------------------------------------------------------------------------- -# We need to be able to import the module under test. The studio backend -# uses relative-style imports (``from utils.…``), so we add the backend -# directory to *sys.path* if it is not already there. +# The studio backend uses relative-style imports (``from utils.…``), so +# add the backend directory to *sys.path* if not already present. # --------------------------------------------------------------------------- import sys @@ -20,8 +19,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Stub the custom logger before importing the module under test so it -# doesn't fail on the ``from loggers import get_logger`` line. +# Stub the custom logger before import so ``from loggers import +# get_logger`` doesn't fail. import types as _types _loggers_stub = _types.ModuleType("loggers") @@ -31,8 +30,11 @@ sys.modules.setdefault("loggers", _loggers_stub) from utils.transformers_version import ( _resolve_base_model, _check_tokenizer_config_needs_v5, + _check_config_needs_510, _check_config_needs_550, + _config_json_cache, _tokenizer_class_cache, + _config_needs_510_cache, _config_needs_550_cache, needs_transformers_5, get_transformers_tier, @@ -90,7 +92,7 @@ class TestResolveBaseModel: (tmp_path / "config.json").write_text(json.dumps(config_cfg)) result = _resolve_base_model(str(tmp_path)) - # Should fall through, not return the self-referencing path + # Falls through; does not return the self-referencing path. assert result == str(tmp_path) def test_no_config_files(self, tmp_path: Path): @@ -174,7 +176,7 @@ class TestNeedsTransformers5: def test_llama_does_not_need_v5(self): """Standard models should not trigger v5.""" - # Patch network call to avoid real fetch + # Patch network call to avoid a real fetch. with patch( "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False, @@ -182,13 +184,12 @@ class TestNeedsTransformers5: assert needs_transformers_5("meta-llama/Llama-3-8B") is False def test_local_checkpoint_resolved_via_config(self, tmp_path: Path): - """A local checkpoint with config.json pointing to Qwen3.5 should need v5.""" + """Local checkpoint with config.json pointing to Qwen3.5 needs v5.""" config_cfg = {"model_name": "Qwen/Qwen3.5-9B"} (tmp_path / "config.json").write_text(json.dumps(config_cfg)) - # _resolve_base_model is called by ensure_transformers_version, - # but needs_transformers_5 just does substring matching. - # We test the full resolution chain here: + # needs_transformers_5 only does substring matching, so test the + # full resolution chain via _resolve_base_model here. resolved = _resolve_base_model(str(tmp_path)) assert needs_transformers_5(resolved) is True @@ -202,6 +203,7 @@ class TestCheckConfigNeeds550: """Tests for _check_config_needs_550() local config.json checks.""" def setup_method(self): + _config_json_cache.clear() _config_needs_550_cache.clear() def test_gemma4_architecture(self, tmp_path: Path): @@ -230,7 +232,7 @@ class TestCheckConfigNeeds550: def test_no_config_json(self, tmp_path: Path): """Missing config.json should return False (fail-open).""" - # Patch network call to avoid real fetch + # Patch network call to avoid a real fetch. with patch("urllib.request.urlopen") as mock_urlopen: mock_urlopen.side_effect = Exception("no network") assert _check_config_needs_550(str(tmp_path)) is False @@ -255,6 +257,106 @@ class TestCheckConfigNeeds550: mock_urlopen.assert_not_called() +# --------------------------------------------------------------------------- +# _check_config_needs_510 — config.json architecture/model_type check +# --------------------------------------------------------------------------- + + +class TestCheckConfigNeeds510: + """Tests for _check_config_needs_510() local config.json checks.""" + + def setup_method(self): + _config_json_cache.clear() + _config_needs_510_cache.clear() + + def test_gemma4_unified_architecture(self, tmp_path: Path): + """config.json with Gemma4UnifiedForConditionalGeneration should return True.""" + cfg = { + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "model_type": "gemma4_unified", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_510(str(tmp_path)) is True + + def test_gemma4_unified_model_type_only(self, tmp_path: Path): + """config.json with model_type=gemma4_unified should return True.""" + cfg = {"model_type": "gemma4_unified"} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_510(str(tmp_path)) is True + + def test_gemma4_unified_assistant_architecture(self, tmp_path: Path): + """Assistant Gemma 4 Unified configs should return True.""" + cfg = { + "architectures": ["Gemma4UnifiedAssistantForCausalLM"], + "model_type": "gemma4_unified_assistant", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_510(str(tmp_path)) is True + + def test_gemma4_unified_assistant_model_type_only(self, tmp_path: Path): + """Assistant Gemma 4 Unified model_type should return True.""" + cfg = {"model_type": "gemma4_unified_assistant"} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_510(str(tmp_path)) is True + + def test_gemma4_assistant_architecture(self, tmp_path: Path): + """Assistant Gemma 4 configs should return True.""" + cfg = { + "architectures": ["Gemma4AssistantForCausalLM"], + "model_type": "gemma4_assistant", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_510(str(tmp_path)) is True + + def test_gemma4_assistant_model_type_only(self, tmp_path: Path): + """Assistant Gemma 4 model_type should return True.""" + cfg = {"model_type": "gemma4_assistant"} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_510(str(tmp_path)) is True + + def test_gemma4_non_unified_returns_false(self, tmp_path: Path): + """Older Gemma 4 config should stay on the 550 tier.""" + cfg = { + "architectures": ["Gemma4ForConditionalGeneration"], + "model_type": "gemma4", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_510(str(tmp_path)) is False + + def test_no_config_json(self, tmp_path: Path): + """Missing config.json should return False (fail-open).""" + # Patch network call to avoid real fetch + with patch("urllib.request.urlopen") as mock_urlopen: + mock_urlopen.side_effect = Exception("no network") + assert _check_config_needs_510(str(tmp_path)) is False + + def test_result_is_cached(self, tmp_path: Path): + """Subsequent calls should use the cache.""" + cfg = {"architectures": ["Gemma4UnifiedForConditionalGeneration"]} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + key = str(tmp_path) + _check_config_needs_510(key) + assert key in _config_needs_510_cache + assert _config_needs_510_cache[key] is True + + def test_local_file_skips_network(self, tmp_path: Path): + """When local config.json exists, no network request should be made.""" + cfg = {"architectures": ["LlamaForCausalLM"]} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + with patch("urllib.request.urlopen") as mock_urlopen: + _check_config_needs_510(str(tmp_path)) + mock_urlopen.assert_not_called() + + # --------------------------------------------------------------------------- # get_transformers_tier — tier detection # --------------------------------------------------------------------------- @@ -265,11 +367,19 @@ class TestGetTransformersTier: def setup_method(self): _tokenizer_class_cache.clear() + _config_json_cache.clear() + _config_needs_510_cache.clear() _config_needs_550_cache.clear() def test_gemma4_substring_returns_550(self): assert get_transformers_tier("google/gemma-4-E2B-it") == "550" + def test_gemma4_12b_substring_returns_510(self): + assert get_transformers_tier("unsloth/gemma-4-12b-it") == "510" + + def test_gemma4_assistant_substring_returns_510(self): + assert get_transformers_tier("google/gemma-4-E2B-it-assistant") == "510" + def test_gemma4_alt_substring_returns_550(self): assert get_transformers_tier("unsloth/gemma4-E4B-it") == "550" @@ -283,21 +393,94 @@ class TestGetTransformersTier: assert get_transformers_tier(str(tmp_path)) == "550" + def test_gemma4_unified_config_json_returns_510(self, tmp_path: Path): + """Local checkpoint with Gemma4 Unified architecture → 510.""" + cfg = { + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "model_type": "gemma4_unified", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert get_transformers_tier(str(tmp_path)) == "510" + + def test_gemma4_assistant_config_json_returns_510(self, tmp_path: Path): + """Local checkpoint with Gemma4 Assistant architecture → 510.""" + cfg = { + "architectures": ["Gemma4AssistantForCausalLM"], + "model_type": "gemma4_assistant", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert get_transformers_tier(str(tmp_path)) == "510" + + def test_local_config_json_short_circuits_path_substrings(self, tmp_path: Path): + """Local config.json should prevent false matches from parent directory names.""" + model_dir = tmp_path / "gemma-4-12b-experiment" / "llama-checkpoint" + model_dir.mkdir(parents = True) + (model_dir / "config.json").write_text( + json.dumps( + { + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + } + ) + ) + (model_dir / "tokenizer_config.json").write_text( + json.dumps({"tokenizer_class": "LlamaTokenizerFast"}) + ) + + with patch("urllib.request.urlopen") as mock_urlopen: + assert get_transformers_tier(str(model_dir)) == "default" + mock_urlopen.assert_not_called() + + def test_remote_config_json_is_fetched_once_for_config_tiers(self): + """510 and 550 slow-path checks should share one config.json fetch.""" + + class _Response: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def read(self): + return json.dumps( + { + "architectures": ["Gemma4ForConditionalGeneration"], + "model_type": "gemma4", + } + ).encode() + + with patch("urllib.request.urlopen", return_value = _Response()) as mock_urlopen: + assert get_transformers_tier("org/no-fast-substring-model") == "550" + + assert mock_urlopen.call_count == 1 + def test_qwen35_returns_530(self): - with patch( - "utils.transformers_version._check_config_needs_550", - return_value = False, + with ( + patch( + "utils.transformers_version._check_config_needs_550", + return_value = False, + ), + patch( + "utils.transformers_version._check_config_needs_510", + return_value = False, + ), ): assert get_transformers_tier("Qwen/Qwen3.5-9B") == "530" def test_ministral_returns_530(self): - with patch( - "utils.transformers_version._check_config_needs_550", - return_value = False, + with ( + patch( + "utils.transformers_version._check_config_needs_550", + return_value = False, + ), + patch( + "utils.transformers_version._check_config_needs_510", + return_value = False, + ), ): - assert ( - get_transformers_tier("mistralai/Ministral-3-8B-Instruct-2512") == "530" - ) + assert get_transformers_tier("mistralai/Ministral-3-8B-Instruct-2512") == "530" def test_llama_returns_default(self): with ( @@ -305,6 +488,10 @@ class TestGetTransformersTier: "utils.transformers_version._check_config_needs_550", return_value = False, ), + patch( + "utils.transformers_version._check_config_needs_510", + return_value = False, + ), patch( "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False, @@ -313,16 +500,22 @@ class TestGetTransformersTier: assert get_transformers_tier("meta-llama/Llama-3-8B") == "default" def test_550_checked_before_530(self): - """Ensure 5.5.0 is checked first — a model matching both should get 550.""" - # This shouldn't happen in practice, but verifies priority + """5.5.0 is checked before 5.3.0 - a model matching both gets 550.""" assert get_transformers_tier("gemma-4-model") == "550" def test_needs_transformers_5_compat(self): - """needs_transformers_5 should return True for both 530 and 550 models.""" + """needs_transformers_5 should return True for 510, 530, and 550 models.""" + assert needs_transformers_5("unsloth/gemma-4-12b-it") is True assert needs_transformers_5("google/gemma-4-E2B-it") is True - with patch( - "utils.transformers_version._check_config_needs_550", - return_value = False, + with ( + patch( + "utils.transformers_version._check_config_needs_550", + return_value = False, + ), + patch( + "utils.transformers_version._check_config_needs_510", + return_value = False, + ), ): assert needs_transformers_5("Qwen/Qwen3.5-9B") is True with ( @@ -330,6 +523,10 @@ class TestGetTransformersTier: "utils.transformers_version._check_config_needs_550", return_value = False, ), + patch( + "utils.transformers_version._check_config_needs_510", + return_value = False, + ), patch( "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False, diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 64c9907119..c66d56528a 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -1,20 +1,10 @@ # 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 utils/hardware and utils/utils — device detection, GPU memory, error formatting. +"""Tests for utils/hardware and utils/utils: device detection, GPU memory, error formatting. -These tests are designed to pass on ANY platform: - • NVIDIA GPU (CUDA backend, requires torch) - • Apple Silicon (MLX backend, requires mlx) - • CPU-only (no GPU at all) - -No ML framework is imported at the top level. -Tests that need torch/mlx internals for mocking are skipped when unavailable. - -Run with: - cd studio/backend - python -m pytest tests/test_utils.py -v +Passes on any platform (NVIDIA/CUDA, Apple Silicon/MLX, CPU-only). No ML framework +is imported at top level; tests needing torch/mlx internals skip when unavailable. """ import platform @@ -25,14 +15,12 @@ import pytest # --- Conditional framework imports --- try: import torch - HAS_TORCH = True except ImportError: HAS_TORCH = False try: import mlx.core as mx - HAS_MLX = True except ImportError: HAS_MLX = False @@ -191,20 +179,15 @@ class TestGetGpuMemoryInfo: assert "backend" in get_gpu_memory_info() def test_backend_matches_device(self): - # The backend field uses _backend_label, which swaps "cuda" for - # "rocm" when running on an AMD host (IS_ROCM=True) so the UI - # can render the correct label. On CUDA / XPU / MLX / CPU hosts - # it is equivalent to `get_device().value`. + # _backend_label swaps "cuda" for "rocm" on AMD hosts; elsewhere it + # equals get_device().value. from utils.hardware.hardware import _backend_label - result = get_gpu_memory_info() assert result["backend"] == _backend_label(get_device()) # --- When a GPU IS available --- - @pytest.mark.skipif( - _actual_device() == "cpu", reason = "No GPU available on this machine" - ) + @pytest.mark.skipif(_actual_device() == "cpu", reason = "No GPU available on this machine") def test_gpu_available_fields(self): result = get_gpu_memory_info() assert result["available"] is True @@ -302,9 +285,7 @@ class TestLogGpuMemory: "free_gb": 14.0, } - with patch( - "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info - ): + with patch("utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info): log_gpu_memory("unit-test") captured = capfd.readouterr() @@ -315,9 +296,7 @@ class TestLogGpuMemory: def test_logs_cpu_fallback_when_no_gpu(self, capfd): fake_info = {"available": False, "backend": "cpu"} - with patch( - "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info - ): + with patch("utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info): log_gpu_memory("cpu-test") captured = capfd.readouterr() diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py index 9e7bbdd1fb..2fee50d842 100644 --- a/studio/backend/tests/test_vision_cache.py +++ b/studio/backend/tests/test_vision_cache.py @@ -3,14 +3,13 @@ """Tests for is_vision_model() caching behaviour. -The vision detection cache (``_vision_detection_cache``) mirrors the existing -``_audio_detection_cache`` pattern used by ``detect_audio_type()``. These -tests verify that: +``_vision_detection_cache`` mirrors the ``_audio_detection_cache`` +pattern used by ``detect_audio_type()``. These tests verify: -* Repeated calls for the same model hit the cache (no redundant work). +* Repeated calls for the same model hit the cache. * Different models each trigger their own detection. * Both True and False results are cached. -* The subprocess path (transformers 5.x models) is also cached. +* The subprocess path (transformers 5.x models) is cached. * Exceptions that fall back to False are cached. """ @@ -21,9 +20,7 @@ from unittest.mock import patch, MagicMock import pytest -# --------------------------------------------------------------------------- # sys.path + logger stub — same pattern as the rest of the test suite -# --------------------------------------------------------------------------- _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) @@ -39,9 +36,7 @@ from utils.models.model_config import ( ) -# --------------------------------------------------------------------------- # Helpers -# --------------------------------------------------------------------------- @pytest.fixture(autouse = True) @@ -52,9 +47,7 @@ def _clear_vision_cache(): _vision_detection_cache.clear() -# --------------------------------------------------------------------------- # Cache hit / miss tests -# --------------------------------------------------------------------------- class TestVisionCacheHitMiss: @@ -62,8 +55,7 @@ class TestVisionCacheHitMiss: @patch("utils.models.model_config._is_vision_model_uncached", return_value = True) def test_second_call_uses_cache(self, mock_uncached): - """Calling is_vision_model() twice for the same model should invoke - the uncached function only once.""" + """Two calls for the same model invoke the uncached fn once.""" assert is_vision_model("org/my-vlm") is True assert is_vision_model("org/my-vlm") is True mock_uncached.assert_called_once_with("org/my-vlm", None) @@ -95,21 +87,19 @@ class TestVisionCacheStoresFalse: assert _vision_detection_cache[("org/text-only", None)] is False -# --------------------------------------------------------------------------- # Subprocess path (transformers 5.x) caching -# --------------------------------------------------------------------------- class TestVisionCacheSubprocessPath: - """Models needing transformers 5.x go through _is_vision_model_subprocess. - The cache should prevent the subprocess from being spawned more than once - per model per process.""" + """transformers 5.x models go through _is_vision_model_subprocess. + The cache should spawn the subprocess at most once per model per + process.""" @patch("utils.models.model_config._is_vision_model_subprocess", return_value = True) @patch("utils.transformers_version.needs_transformers_5", return_value = True) def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess): - """Subprocess should only fire on the first call; second is cached.""" - # First call: goes through uncached → subprocess + """Subprocess fires only on the first call; second is cached.""" + # First call: uncached → subprocess assert is_vision_model("unsloth/Qwen3.5-2B") is True # Second call: cache hit, no subprocess assert is_vision_model("unsloth/Qwen3.5-2B") is True @@ -117,17 +107,26 @@ class TestVisionCacheSubprocessPath: mock_subprocess.assert_called_once() assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None)] is True + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = True) + @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None) + @patch("utils.transformers_version.needs_transformers_5", return_value = True) + def test_subprocess_none_falls_back_to_raw_vision_config( + self, mock_needs_t5, mock_subprocess, mock_raw_config + ): + assert is_vision_model("unsloth/gemma-4-E4B-it") is True + assert is_vision_model("unsloth/gemma-4-E4B-it") is True + + mock_subprocess.assert_called_once() + mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None) + -# --------------------------------------------------------------------------- # Exception handling — cache the False fallback -# --------------------------------------------------------------------------- class TestVisionCacheOnException: - """When detection raises an exception, _is_vision_model_uncached - distinguishes permanent failures (cached as False) from transient - failures (returned as None, not cached so the next call can retry). - Verify both contracts.""" + """On exception, _is_vision_model_uncached distinguishes permanent + failures (cached as False) from transient ones (returned as None, + not cached, so the next call retries). Verify both contracts.""" @patch( "utils.models.model_config.load_model_config", @@ -136,16 +135,11 @@ class TestVisionCacheOnException: @patch("utils.transformers_version.needs_transformers_5", return_value = False) def test_permanent_exception_result_cached(self, mock_needs_t5, mock_load_config): """A permanent failure (ValueError / RepositoryNotFoundError / - GatedRepoError / JSONDecodeError) should be caught, return False, - and that False should be cached so subsequent calls don't retry. - - ValueError is used here because it's the simplest of the - code-path's cacheable exception types and does not require an - import of huggingface_hub errors (whose module path varies - across versions).""" - # First call: load_model_config raises -> except branch -> False. + GatedRepoError / JSONDecodeError) is caught, returns False, and + that False is cached so subsequent calls don't retry. ValueError + stands in as the simplest cacheable exception type.""" + # First call raises -> False; second is a cache hit. assert is_vision_model("broken/model") is False - # Second call: cache hit, load_model_config not called again. assert is_vision_model("broken/model") is False mock_load_config.assert_called_once() @@ -155,28 +149,20 @@ class TestVisionCacheOnException: ) @patch("utils.transformers_version.needs_transformers_5", return_value = False) def test_transient_exception_not_cached(self, mock_needs_t5, mock_load_config): - """A transient failure (OSError, timeouts) should return None from - _is_vision_model_uncached, surface as False to the caller, and - NOT be cached, so the next call retries detection. This matches - the documented behaviour on _vision_detection_cache: - 'transient failures (network errors, timeouts) are NOT cached so - they can be retried.'""" - # First call: load_model_config raises OSError -> uncached None - # -> caller returns False without caching. + """A transient failure (OSError, timeouts) returns None from + _is_vision_model_uncached, surfaces as False, and is NOT cached + so the next call retries.""" + # First call: OSError -> False, not cached; second call retries. assert is_vision_model("broken/model") is False - # Second call: cache miss again, load_model_config called a - # second time. assert is_vision_model("broken/model") is False assert mock_load_config.call_count == 2 -# --------------------------------------------------------------------------- # Direct detection path (non-transformers-5 models) caching -# --------------------------------------------------------------------------- class TestVisionCacheDirectPath: - """For models that do NOT need transformers 5.x, the detection goes through + """Models that do NOT need transformers 5.x detect via load_model_config directly. The cache must work the same way.""" @patch("utils.transformers_version.needs_transformers_5", return_value = False) @@ -202,16 +188,14 @@ class TestVisionCacheDirectPath: cfg.architectures = ["LlamaForCausalLM"] mock_load_config.return_value = cfg - # LlamaForCausalLM doesn't end with VLM suffixes, no vision_config, etc. + # No VLM suffix, no vision_config, etc. assert is_vision_model("meta-llama/Llama-3-8B") is False assert is_vision_model("meta-llama/Llama-3-8B") is False mock_load_config.assert_called_once() @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_vision_config_attr_detected_and_cached( - self, mock_load_config, mock_needs_t5 - ): + def test_vision_config_attr_detected_and_cached(self, mock_load_config, mock_needs_t5): """Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist cfg.model_type = "qwen2_vl" @@ -223,6 +207,42 @@ class TestVisionCacheDirectPath: assert is_vision_model("Qwen/Qwen2-VL-7B") is True mock_load_config.assert_called_once() + @patch("utils.transformers_version.needs_transformers_5", return_value = False) + @patch("utils.models.model_config.load_model_config") + def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5): + cfg = MagicMock(spec = []) + cfg.model_type = "gemma4" + cfg.architectures = ["Gemma4ForConditionalGeneration"] + mock_load_config.return_value = cfg + + assert is_vision_model("google/gemma-4-E4B-it") is True + assert is_vision_model("google/gemma-4-E4B-it") is True + mock_load_config.assert_called_once() + + @patch("utils.transformers_version.needs_transformers_5", return_value = False) + @patch("utils.models.model_config.load_model_config") + def test_gemma4_audio_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5): + cfg = MagicMock(spec = []) + cfg.model_type = "gemma4_audio" + cfg.architectures = ["Gemma4AudioModel"] + mock_load_config.return_value = cfg + + assert is_vision_model("local/gemma4-audio-encoder") is False + assert is_vision_model("local/gemma4-audio-encoder") is False + mock_load_config.assert_called_once() + + @patch("utils.transformers_version.needs_transformers_5", return_value = False) + @patch("utils.models.model_config.load_model_config") + def test_gemma4_text_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5): + cfg = MagicMock(spec = []) + cfg.model_type = "gemma4_text" + cfg.architectures = ["Gemma4ForCausalLM"] + mock_load_config.return_value = cfg + + assert is_vision_model("local/gemma-4-text") is False + assert is_vision_model("local/gemma-4-text") is False + mock_load_config.assert_called_once() + @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5): @@ -238,21 +258,18 @@ class TestVisionCacheDirectPath: mock_load_config.assert_called_once() -# --------------------------------------------------------------------------- # hf_token handling -# --------------------------------------------------------------------------- class TestVisionCacheTokenHandling: - """The cache is keyed on (model_name, hf_token). - Different tokens for the same model should trigger separate detections - to handle gated models correctly.""" + """The cache is keyed on (model_name, hf_token). Different tokens + for the same model trigger separate detections for gated models.""" @patch("utils.models.model_config._is_vision_model_uncached", return_value = True) def test_different_tokens_trigger_new_detection(self, mock_uncached): - """Calls with different tokens should trigger separate detections to - handle gated models correctly (e.g. unauthenticated probe → False, - then authenticated call should re-check).""" + """Different tokens trigger separate detections for gated models + (e.g. unauthenticated probe → False, then authenticated + re-check).""" assert is_vision_model("gated/model", hf_token = "token-a") is True assert is_vision_model("gated/model", hf_token = "token-b") is True assert mock_uncached.call_count == 2 @@ -263,3 +280,157 @@ class TestVisionCacheTokenHandling: assert is_vision_model("gated/model", hf_token = "token-a") is True assert is_vision_model("gated/model", hf_token = "token-a") is True mock_uncached.assert_called_once() + + +# --------------------------------------------------------------------------- +# Direct unit tests for _raw_config_has_vision_config +# --------------------------------------------------------------------------- + + +import json as _json + +from utils.models.model_config import ( + _AUDIO_ONLY_MODEL_TYPES, + _VISION_CHECK_INLINE_HELPERS, + _VISION_CHECK_SCRIPT, + _is_vlm, + _raw_config_has_vision_config, +) + + +def _write_config(tmp_path, config): + (tmp_path / "config.json").write_text(_json.dumps(config)) + return tmp_path + + +class TestRawConfigVlmDetection: + """Direct coverage of _raw_config_has_vision_config across the same + indicator set used by _is_vlm. The cache integration tests above mock + this function; these exercise its real implementation.""" + + def test_truthy_vision_config(self, tmp_path): + p = _write_config(tmp_path, {"vision_config": {"hidden_size": 1024}}) + assert _raw_config_has_vision_config(str(p)) is True + + def test_empty_vision_config_key(self, tmp_path): + p = _write_config(tmp_path, {"vision_config": {}}) + assert _raw_config_has_vision_config(str(p)) is True + + def test_arch_suffix_detection(self, tmp_path): + p = _write_config( + tmp_path, + { + "architectures": ["Gemma4ForConditionalGeneration"], + "model_type": "gemma4", + }, + ) + assert _raw_config_has_vision_config(str(p)) is True + + def test_img_processor_key(self, tmp_path): + p = _write_config(tmp_path, {"img_processor": {"image_size": 336}}) + assert _raw_config_has_vision_config(str(p)) is True + + def test_image_token_index_key(self, tmp_path): + p = _write_config(tmp_path, {"image_token_index": 32000}) + assert _raw_config_has_vision_config(str(p)) is True + + def test_known_vlm_model_type(self, tmp_path): + p = _write_config(tmp_path, {"model_type": "gemma4"}) + assert _raw_config_has_vision_config(str(p)) is True + + def test_plain_text_model_returns_false(self, tmp_path): + p = _write_config( + tmp_path, + {"model_type": "llama", "architectures": ["LlamaForCausalLM"]}, + ) + assert _raw_config_has_vision_config(str(p)) is False + + def test_missing_config_returns_none(self, tmp_path): + assert _raw_config_has_vision_config(str(tmp_path)) is None + + +# --------------------------------------------------------------------------- +# Self-contained subprocess script (no parent backend imports) +# --------------------------------------------------------------------------- + + +class TestSubprocessScript: + def test_does_not_import_parent_module(self): + assert "from utils.models.model_config" not in _VISION_CHECK_SCRIPT + + def test_inline_is_vlm_executes_correctly(self): + ns: dict = {} + exec(_VISION_CHECK_INLINE_HELPERS, ns) + inline_is_vlm = ns["_is_vlm"] + + class _C: + def __init__(self, **kw): + for k, v in kw.items(): + setattr(self, k, v) + + assert ( + inline_is_vlm( + _C( + model_type = "gemma4", + architectures = ["Gemma4ForConditionalGeneration"], + ) + ) + is True + ) + assert ( + inline_is_vlm(_C(model_type = "gemma4_text", architectures = ["Gemma4ForCausalLM"])) + is False + ) + assert inline_is_vlm(_C(model_type = "llama", architectures = ["LlamaForCausalLM"])) is False + + +# --------------------------------------------------------------------------- +# Audio-only model exclusion must apply across every detection path +# --------------------------------------------------------------------------- + + +class TestVlmAudioExclusion: + """The {csm, whisper} guard previously lived only in the direct caller + branch. These tests assert it now applies inside _is_vlm, the raw + fallback, and the inlined subprocess helper too.""" + + def test_audio_only_set_canonical(self): + assert _AUDIO_ONLY_MODEL_TYPES == {"csm", "whisper"} + + def test_is_vlm_excludes_whisper(self): + cfg = MagicMock(spec = []) + cfg.model_type = "whisper" + cfg.architectures = ["WhisperForConditionalGeneration"] + assert _is_vlm(cfg) is False + + def test_raw_fallback_excludes_whisper(self, tmp_path): + p = _write_config( + tmp_path, + { + "architectures": ["WhisperForConditionalGeneration"], + "model_type": "whisper", + }, + ) + assert _raw_config_has_vision_config(str(p)) is False + + def test_inline_subprocess_helper_excludes_whisper(self): + ns: dict = {} + exec(_VISION_CHECK_INLINE_HELPERS, ns) + cfg = MagicMock(spec = []) + cfg.model_type = "whisper" + cfg.architectures = ["WhisperForConditionalGeneration"] + assert ns["_is_vlm"](cfg) is False + + @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None) + @patch("utils.transformers_version.needs_transformers_5", return_value = True) + def test_t5_subprocess_none_falls_back_through_raw_for_whisper( + self, mock_needs_t5, mock_subprocess, tmp_path + ): + _write_config( + tmp_path, + { + "architectures": ["WhisperForConditionalGeneration"], + "model_type": "whisper", + }, + ) + assert is_vision_model(str(tmp_path)) is False diff --git a/studio/backend/tests/test_vram_estimation.py b/studio/backend/tests/test_vram_estimation.py index e54ae6dcf8..2def8738e2 100644 --- a/studio/backend/tests/test_vram_estimation.py +++ b/studio/backend/tests/test_vram_estimation.py @@ -316,12 +316,8 @@ class TestLoraParams(unittest.TestCase): self.assertLess(qv_only, all_mods) def test_moe_mlp_modules_scale_with_experts(self): - dense_lora = compute_lora_params( - LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"] - ) - moe_lora = compute_lora_params( - MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"] - ) + dense_lora = compute_lora_params(LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"]) + moe_lora = compute_lora_params(MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"]) ratio = moe_lora / dense_lora self.assertAlmostEqual(ratio, 8.0, delta = 0.5) @@ -338,12 +334,8 @@ class TestLoraParams(unittest.TestCase): self.assertGreater(moe_lora, dense_lora * 20) def test_attention_modules_same_for_moe(self): - dense_attn = compute_lora_params( - LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] - ) - moe_attn = compute_lora_params( - MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] - ) + dense_attn = compute_lora_params(LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]) + moe_attn = compute_lora_params(MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]) self.assertEqual(dense_attn, moe_attn) def test_all_linear_uses_default_text_modules(self): @@ -466,9 +458,7 @@ class TestActivationBytes(unittest.TestCase): def test_non_flash_attention_uses_quadratic_path(self): seq_len = 4096 - expected_quadratic = ( - 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 - ) + expected_quadratic = 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 for attention_implementation in ("eager", "unknown_impl", None): with self.subTest(attention_implementation = attention_implementation): non_flash = compute_activation_bytes( @@ -483,9 +473,7 @@ class TestActivationBytes(unittest.TestCase): def test_non_flash_attention_without_gc_scales_quadratic_path_by_layers(self): seq_len = 4096 - one_layer = ( - 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 - ) + one_layer = 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 non_flash = compute_activation_bytes( STRUCTURED_MIXED, 1, @@ -543,7 +531,7 @@ class TestQuantizationSkips(unittest.TestCase): ) def test_vlm_prefix_skip_module_does_not_match_text_alias(self): - # vision_tower-prefixed skips must not shadow text aliases sharing the + # vision_tower-prefixed skips must not shadow text aliases with the # same suffix. baseline = replace(QUANT_SKIP_STRUCTURED, quantization_skip_modules = []) vlm_skip = replace( @@ -717,9 +705,7 @@ class TestEstimateTrainingVram(unittest.TestCase): ) v8 = estimate_training_vram(LLAMA_8B, opt8) v32 = estimate_training_vram(LLAMA_8B, opt32) - self.assertAlmostEqual( - v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1 - ) + self.assertAlmostEqual(v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1) def test_min_gpu_vram_treats_activations_as_per_gpu_fixed(self): config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) @@ -769,9 +755,7 @@ class TestEstimateTrainingVram(unittest.TestCase): optimizer = "adamw_8bit", load_in_4bit = False, ) - expected_floor = int( - compute_model_weights_bytes(LLAMA_8B, "full", False) * 0.15 - ) + expected_floor = int(compute_model_weights_bytes(LLAMA_8B, "full", False) * 0.15) with patch( "utils.hardware.vram_estimation.compute_gradient_bytes", return_value = 1, @@ -1011,9 +995,9 @@ class TestParallelDenseMoE(unittest.TestCase): + with_parallel.num_experts * with_parallel.hidden_size ) dense_only = with_parallel.hidden_size * with_parallel.intermediate_size * 3 - # why: under gemma4 enable_moe_block, the layer's `self.experts` is a - # sibling of `self.mlp`; the `text.layers..mlp` aggregate must - # cover the dense path only, with experts in their own aggregate. + # why: under gemma4 enable_moe_block, `self.experts` is a sibling of + # `self.mlp`; the `text.layers..mlp` aggregate covers the dense path + # only, with experts in their own aggregate. self.assertEqual(elements["text.layers.0.mlp"], dense_only) self.assertEqual(elements["text.layers.0.experts"], moe_only) @@ -1059,7 +1043,6 @@ class TestDenseLayerIndices(unittest.TestCase): class TestKvSharedLayer(unittest.TestCase): def test_fully_shared_kv_returns_false_matching_upstream(self): from utils.hardware.vram_estimation import _is_kv_shared_layer - arch = ModelArchConfig( hidden_size = 512, num_hidden_layers = 4, @@ -1165,9 +1148,9 @@ class TestPerLayerInputAccounting(unittest.TestCase): def test_per_layer_input_modules_count_quantizable_block(self): with_ple = self._arch() without_ple = replace(with_ple, hidden_size_per_layer_input = 0) - # The PLE block adds: model_projection (hd*nl*pli), per_layer_input_gate - # (hd*pli per layer) + per_layer_projection (pli*hd per layer) as - # quantizable text linears. + # PLE block adds these quantizable text linears: model_projection + # (hd*nl*pli), per_layer_input_gate (hd*pli per layer), + # per_layer_projection (pli*hd per layer). n_layers = with_ple.num_hidden_layers hd = with_ple.hidden_size pli = with_ple.hidden_size_per_layer_input @@ -1178,10 +1161,10 @@ class TestPerLayerInputAccounting(unittest.TestCase): self.assertGreaterEqual(delta, expected_quantizable_extra) def test_all_linear_lora_excludes_per_layer_input_modules(self): - # why: Unsloth's get_peft_regex requires module names to contain a - # component tag (mlp/attn/...); PLE module names (per_layer_input_gate, - # per_layer_projection, per_layer_model_projection) lack any tag, so - # all-linear training does NOT attach LoRA to them. + # why: Unsloth's get_peft_regex requires a component tag (mlp/attn/...) + # in module names; PLE names (per_layer_input_gate, per_layer_projection, + # per_layer_model_projection) lack one, so all-linear does NOT attach + # LoRA to them. arch = self._arch() without_ple = replace(arch, hidden_size_per_layer_input = 0) self.assertEqual( @@ -1251,11 +1234,10 @@ class TestExpertsSkipGranularity(unittest.TestCase): bytes_skip_experts = compute_model_weights_bytes(skip_experts, "qlora", True) bytes_skip_mlp = compute_model_weights_bytes(skip_full_mlp, "qlora", True) # why: under gemma4 enable_moe_block, `self.experts` is a sibling of - # `self.mlp`; skipping `model.layers.0.mlp` should cover only the - # dense MLP, while `model.layers.0.mlp.experts` covers the routed - # experts. Routed experts have far more params than the dense MLP, - # so skipping experts must add more bytes than skipping the dense - # path. + # `self.mlp`; skipping `model.layers.0.mlp` covers only the dense MLP, + # while `model.layers.0.mlp.experts` covers the routed experts. Routed + # experts have far more params than the dense MLP, so skipping experts + # must add more bytes than skipping the dense path. self.assertGreater(bytes_skip_experts, bytes_no_skip) self.assertGreater(bytes_skip_mlp, bytes_no_skip) self.assertGreater(bytes_skip_experts, bytes_skip_mlp) @@ -1293,9 +1275,7 @@ class TestSharedExperts(unittest.TestCase): delta_per_layer = 4096 * 1407 * 3 * 2 expected_delta = delta_per_layer * 32 * 2 actual_delta = w_yes - w_no - self.assertAlmostEqual( - actual_delta, expected_delta, delta = expected_delta * 0.01 - ) + self.assertAlmostEqual(actual_delta, expected_delta, delta = expected_delta * 0.01) def test_deepseek_v3_params_in_range(self): total = compute_total_params(DEEPSEEK_V3) @@ -1411,9 +1391,7 @@ class TestDenseMoEMix(unittest.TestCase): moe_intermediate_size = 1024, num_dense_layers = 5, ) - lora_all = compute_lora_params( - all_moe, 16, ["gate_proj", "up_proj", "down_proj"] - ) + lora_all = compute_lora_params(all_moe, 16, ["gate_proj", "up_proj", "down_proj"]) lora_mix = compute_lora_params(mixed, 16, ["gate_proj", "up_proj", "down_proj"]) self.assertNotEqual(lora_all, lora_mix) @@ -1497,9 +1475,7 @@ class TestPerLayerInputSkipAlias(unittest.TestCase): delta = _compute_skipped_quantizable_elements(arch) self.assertEqual( delta, - arch.hidden_size - * arch.num_hidden_layers - * arch.hidden_size_per_layer_input, + arch.hidden_size * arch.num_hidden_layers * arch.hidden_size_per_layer_input, ) def test_layer_aggregate_skip_includes_per_layer_input_modules(self): @@ -1508,8 +1484,8 @@ class TestPerLayerInputSkipAlias(unittest.TestCase): ) arch_with = extract_arch_config(self._hf(["model.layers.0"])) - # The text.layers.0 aggregate must include the PLE per-layer modules, - # so the same skip on a config without PLE produces a smaller value. + # text.layers.0 aggregate includes the PLE per-layer modules, so the + # same skip on a no-PLE config produces a smaller value. arch_without = extract_arch_config( SimpleNamespace( text_config = SimpleNamespace( @@ -1578,12 +1554,10 @@ class TestSharedExpertVariants(unittest.TestCase): def test_shared_expert_size_separate_from_routed_changes_weight_count(self): from utils.hardware.vram_estimation import _compute_moe_mlp_elements - arch_separate = extract_arch_config( - self._hf(shared_expert_intermediate_size = 64) - ) + arch_separate = extract_arch_config(self._hf(shared_expert_intermediate_size = 64)) arch_implicit = extract_arch_config(self._hf(n_shared_experts = 1)) # Different shared sizes (64 vs default moe_intermediate_size=128) must - # produce different MoE element counts. + # give different MoE element counts. self.assertNotEqual( _compute_moe_mlp_elements(arch_separate), _compute_moe_mlp_elements(arch_implicit), @@ -1592,7 +1566,7 @@ class TestSharedExpertVariants(unittest.TestCase): def test_shared_expert_gate_counted_only_for_qwen_style(self): from utils.hardware.vram_estimation import _compute_moe_mlp_elements - # Qwen-style: shared_expert_intermediate_size set -> shared_expert_gate counted. + # Qwen-style: shared_expert_intermediate_size set -> gate counted. qwen_arch = extract_arch_config(self._hf(shared_expert_intermediate_size = 64)) hd = qwen_arch.hidden_size ms = qwen_arch.moe_intermediate_size @@ -1624,9 +1598,7 @@ class TestSharedExpertActivation(unittest.TestCase): moe_intermediate_size = 64, **fields, ) - return extract_arch_config( - SimpleNamespace(text_config = text_config, quantization_config = {}) - ) + return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) def test_shared_expert_increases_activation_bytes(self): with_shared = self._make(shared_expert_intermediate_size = 64) @@ -1651,8 +1623,8 @@ class TestSharedExpertActivation(unittest.TestCase): ) def test_shared_expert_plus_dense_block_compose(self): - # gemma4 enable_moe_block with hypothetical shared expert: dense + routed - # + shared all live per layer; mlp_size should sum all three terms. + # gemma4 enable_moe_block with a hypothetical shared expert: dense + + # routed + shared all live per layer; mlp_size sums all three. from utils.hardware.vram_estimation import _layer_qkv_mlp_sizes arch = self._make( @@ -1678,9 +1650,7 @@ class TestPerLayerInputActivation(unittest.TestCase): tie_word_embeddings = False, **fields, ) - return extract_arch_config( - SimpleNamespace(text_config = text_config, quantization_config = {}) - ) + return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) def test_ple_increases_activation_bytes(self): with_ple = self._make( @@ -1744,9 +1714,7 @@ class TestKvSharedActivation(unittest.TestCase): num_kv_shared_layers = kv_shared, layer_types = ["full_attention"] * 4, ) - return extract_arch_config( - SimpleNamespace(text_config = text_config, quantization_config = {}) - ) + return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) def test_kv_shared_layers_keep_activation_bytes(self): shared = self._make(kv_shared = 2) @@ -1792,10 +1760,7 @@ class TestSparseMoeSkipAliases(unittest.TestCase): def test_gemma4_layers_experts_alias_pulls_routed(self): from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements - - arch = extract_arch_config( - self._hf(["model.layers.0.experts"], enable_moe_block = True) - ) + arch = extract_arch_config(self._hf(["model.layers.0.experts"], enable_moe_block = True)) self.assertGreater(_compute_skipped_quantizable_elements(arch), 0) def test_qwen_shared_expert_skip_pulls_only_shared(self): @@ -1807,7 +1772,7 @@ class TestSparseMoeSkipAliases(unittest.TestCase): shared_expert_intermediate_size = 32, ) ) - # shared_expert delta only -- routed mlp.experts is NOT skipped. + # shared_expert delta only -- routed mlp.experts NOT skipped. delta = _compute_skipped_quantizable_elements(arch) self.assertGreater(delta, 0) full_layer = extract_arch_config( @@ -1823,7 +1788,6 @@ class TestSparseMoeSkipAliases(unittest.TestCase): def test_exaone_shared_experts_plural_alias(self): from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements - arch = extract_arch_config( self._hf( ["model.layers.0.mlp.shared_experts"], @@ -1847,9 +1811,7 @@ class TestAllLinearMoELoraExclusion(unittest.TestCase): moe_intermediate_size = 64, **fields, ) - return extract_arch_config( - SimpleNamespace(text_config = text_config, quantization_config = {}) - ) + return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) def test_all_linear_drops_routed_moe_expert_lora(self): arch = self._arch() @@ -1867,9 +1829,7 @@ class TestAllLinearMoELoraExclusion(unittest.TestCase): def test_all_linear_includes_attention_lora(self): arch = self._arch() all_linear = compute_lora_params(arch, 8, "all-linear") - attn_only = compute_lora_params( - arch, 8, ["q_proj", "k_proj", "v_proj", "o_proj"] - ) + attn_only = compute_lora_params(arch, 8, ["q_proj", "k_proj", "v_proj", "o_proj"]) # all-linear still attaches to attention nn.Linear modules. self.assertGreaterEqual(all_linear, attn_only) @@ -1887,9 +1847,7 @@ class TestExplicitPerLayerInputLora(unittest.TestCase): hidden_size_per_layer_input = 32, vocab_size_per_layer_input = 128, ) - return extract_arch_config( - SimpleNamespace(text_config = text_config, quantization_config = {}) - ) + return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) def test_explicit_per_layer_input_gate_returns_nonzero(self): arch = self._arch() @@ -1928,9 +1886,7 @@ class TestTopKExpertActivation(unittest.TestCase): moe_intermediate_size = 64, **fields, ) - return extract_arch_config( - SimpleNamespace(text_config = text_config, quantization_config = {}) - ) + return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) def test_num_experts_per_tok_extracted(self): arch = self._make(num_experts_per_tok = 4) @@ -1987,10 +1943,10 @@ class TestErnieMoEListConfig(unittest.TestCase): moe_intermediate_size = [1536, 512], ) ) - # why: ERNIE 4.5 VL MoE encodes [text_routed, vision_routed]; the - # second element is the vision-routed expert width, not the shared - # expert width. Shared experts are sized from the text-routed width - # (= moe_intermediate_size[0]) when moe_num_shared_experts is set. + # why: ERNIE 4.5 VL MoE encodes [text_routed, vision_routed]; element 1 + # is the vision-routed width, not the shared-expert width. Shared + # experts size from the text-routed width (moe_intermediate_size[0]) + # when moe_num_shared_experts is set. self.assertEqual(arch.moe_intermediate_size, 1536) self.assertIsNone(arch.shared_expert_intermediate_size) self.assertEqual(arch.n_shared_experts, 0) @@ -2113,7 +2069,7 @@ class TestMultimodalFullModelBytes(unittest.TestCase): load_in_4bit = True, ) self.assertEqual(metadata.get("estimation_mode"), "detailed") - # model_weights_gb must reflect the extra non-text bytes (>5 GB + # model_weights_gb must reflect the extra non-text bytes (>5 GB, # since text-only arch_fp16 is small for these dims). self.assertGreater(metadata["vram_breakdown"]["model_weights_gb"], 5.0) @@ -2180,13 +2136,11 @@ class TestLlama4ArchExtraction(unittest.TestCase): def test_llama4_moe_layers_dispatch_uses_explicit_indices(self): from utils.hardware.vram_estimation import _compute_dense_layer_indices - cfg = SimpleNamespace(num_hidden_layers = 4, moe_layers = [1, 3]) self.assertEqual(_compute_dense_layer_indices(cfg, 4), (0, 2)) def test_llama4_moe_layers_takes_priority_over_first_k_dense_replace(self): from utils.hardware.vram_estimation import _compute_dense_layer_indices - cfg = SimpleNamespace( num_hidden_layers = 6, moe_layers = [2, 4], @@ -2288,7 +2242,6 @@ class TestDbrxFfnConfigExtraction(unittest.TestCase): class TestErniePhaseModuloDispatch(unittest.TestCase): def test_phase_modulo_with_interval_two_matches_decoder(self): from utils.hardware.vram_estimation import _compute_dense_layer_indices - cfg = SimpleNamespace( num_hidden_layers = 10, moe_layer_start_index = 2, @@ -2300,7 +2253,6 @@ class TestErniePhaseModuloDispatch(unittest.TestCase): def test_phase_modulo_with_interval_three(self): from utils.hardware.vram_estimation import _compute_dense_layer_indices - cfg = SimpleNamespace( num_hidden_layers = 9, moe_layer_start_index = 0, diff --git a/studio/backend/tests/test_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py index 023630fb9a..88a1a28d14 100644 --- a/studio/backend/tests/test_windows_gpu_detection_mock.py +++ b/studio/backend/tests/test_windows_gpu_detection_mock.py @@ -3,12 +3,12 @@ """Windows GPU-detection regression test on a synthetic layout. -The bug (#5106): on Windows without a system CUDA toolkit, the prebuilt -llama-server.exe could not LoadLibrary cudart64_X / cublas64_X / +Bug (#5106): on Windows without a system CUDA toolkit, the prebuilt +llama-server.exe couldn't LoadLibrary cudart64_X / cublas64_X / cublasLt64_X, so ggml-cuda.dll's static import on cublas64_X.dll failed and the model fell back to CPU even when nvidia-smi reported the GPU. -The fix: +Fix: * #5322 overlays upstream's paired cudart bundle into install_dir/build/bin/Release/ next to llama-server.exe. * #5324 prepends pip-installed nvidia//{bin,bin/x86_64,Library/ @@ -34,13 +34,9 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Stub heavy deps only if they actually fail to import -- unconditional -# stubs would shadow the real module for sibling tests in this dir. -# Use try-import rather than find_spec: loggers/__init__.py re-exports -# handlers.get_logger, which does `from fastapi import Request, -# Response` at module load. find_spec("loggers") returns a spec even -# without fastapi, but the import then raises. CI has fastapi, so this -# is dev-machine ergonomics only. +# Stub heavy deps only if they fail to import (unconditional stubs would shadow +# the real module for sibling tests). Use try-import, not find_spec: loggers +# imports fastapi at load, so find_spec succeeds but the import then raises. import importlib as _importlib # noqa: E402 @@ -100,7 +96,7 @@ from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 # Upstream b9103 cudart bundle: exactly these three DLLs per CUDA major, -# no executables, no subdirectories. Verified by direct unzip. +# no executables or subdirectories. Verified by direct unzip. REAL_UPSTREAM_CUDART_BUNDLE = { "12.4": ("cudart64_12.dll", "cublas64_12.dll", "cublasLt64_12.dll"), "13.1": ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"), @@ -132,24 +128,21 @@ REAL_PIP_NVIDIA_WHEEL_LAYOUTS = { def _populate_studio_venv(prefix: Path) -> None: - """Lay out fake nvidia + torch wheels in /Lib/site-packages - matching the real win_amd64 wheel layouts. Contents are stub bytes; - only directory structure matters.""" + """Lay out fake nvidia + torch wheels matching real win_amd64 layouts (stub bytes).""" site = prefix / "Lib" / "site-packages" for rel, dlls in REAL_PIP_NVIDIA_WHEEL_LAYOUTS.items(): d = site / Path(rel) d.mkdir(parents = True, exist_ok = True) for name in dlls: (d / name).write_bytes(b"PE-stub") - # install_python_stack always installs torch alongside nvidia. + # install_python_stack always installs torch beside nvidia. (site / "torch" / "lib").mkdir(parents = True, exist_ok = True) for fn in ("c10.dll", "torch.dll", "torch_cpu.dll", "torch_python.dll"): (site / "torch" / "lib" / fn).write_bytes(b"PE-stub") def _populate_studio_install(install_dir: Path, runtime: str = "13.1") -> None: - """Lay out install_dir/build/bin/Release/ as #5322 leaves it: main - archive payload + paired cudart bundle overlay.""" + """Lay out install_dir/build/bin/Release/ as #5322 leaves it: payload + cudart overlay.""" rel = install_dir / "build" / "bin" / "Release" rel.mkdir(parents = True, exist_ok = True) for fn in ( @@ -163,25 +156,23 @@ def _populate_studio_install(install_dir: Path, runtime: str = "13.1") -> None: "mtmd.dll", ): (rel / fn).write_bytes(b"PE-stub") - # The cudart overlay #5322 contributes. + # The cudart overlay from #5322. for fn in REAL_UPSTREAM_CUDART_BUNDLE[runtime]: (rel / fn).write_bytes(b"PE-stub") def _build_path_dirs_like_start_llama_server( - binary_dir: Path, prefix: Path, cuda_path: str = "" + binary_dir: Path, + prefix: Path, + cuda_path: str = "", ) -> list[str]: - """Path-friendly wrapper around LlamaCppBackend._build_windows_path_dirs. - Asserting against the staticmethod (not a hand-copy) is the point: - if the win32 PATH order drops _windows_pip_nvidia_dll_dirs, tests fail.""" - return LlamaCppBackend._build_windows_path_dirs( - str(binary_dir), str(prefix), cuda_path - ) + """Wrapper around the real _build_windows_path_dirs staticmethod.""" + return LlamaCppBackend._build_windows_path_dirs(str(binary_dir), str(prefix), cuda_path) def _mock_nvidia_smi_run(fake_output: str, returncode: int = 0) -> "mock._patch": """Patch subprocess.run so the nvidia-smi probe returns fake_output; - other subprocess.run calls pass through.""" + other calls pass through.""" real_run = subprocess.run def fake_run(cmd, *args, **kwargs): @@ -199,20 +190,18 @@ def _mock_nvidia_smi_run(fake_output: str, returncode: int = 0) -> "mock._patch" # --------------------------------------------------------------------- # class TestWindowsGpuDetectionAfter5106Fix: """End-to-end #5106 fix on a synthetic Windows layout. nvidia-smi - mocked; resolver, PATH builder and install layout exercised live.""" + mocked; resolver, PATH builder, and install layout run live.""" def test_nvidia_smi_probe_reports_synthetic_gpu(self, monkeypatch): """Probe parses CSV output and returns (index, free_mib).""" - # Clear inherited masks so the synthetic CSV is not filtered. + # Clear inherited masks so the synthetic CSV isn't filtered. monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) monkeypatch.delenv("NVIDIA_VISIBLE_DEVICES", raising = False) # The #5106 reporter's exact reproducer: RTX 4090, 22805 MiB. fake_csv = "0, 22805\n" with _mock_nvidia_smi_run(fake_csv): gpus = LlamaCppBackend._get_gpu_free_memory() - assert gpus == [ - (0, 22805) - ], f"GPU probe failed to parse mocked nvidia-smi output: {gpus}" + assert gpus == [(0, 22805)], f"GPU probe failed to parse mocked nvidia-smi output: {gpus}" def test_nvidia_smi_probe_respects_cuda_visible_devices(self, monkeypatch): """CUDA_VISIBLE_DEVICES=1 -> only GPU 1 visible.""" @@ -224,7 +213,7 @@ class TestWindowsGpuDetectionAfter5106Fix: def test_windows_install_dir_has_all_three_cudart_dlls(self, tmp_path): """All three bundle DLLs must land in install_dir/build/bin/ - Release; missing any one breaks ggml-cuda.dll's PE import chain.""" + Release; any missing one breaks ggml-cuda.dll's PE import chain.""" install = tmp_path / "studio_install" _populate_studio_install(install, runtime = "13.1") rel = install / "build" / "bin" / "Release" @@ -234,8 +223,8 @@ class TestWindowsGpuDetectionAfter5106Fix: assert (rel / "ggml-cuda.dll").exists() def test_resolver_finds_real_pypi_wheel_layouts(self, tmp_path): - """Resolver must pick up every real-world wheel layout: - nvidia//bin, nvidia//bin/x86_64, torch/lib.""" + """Resolver must pick up every wheel layout: nvidia//bin, + nvidia//bin/x86_64, torch/lib.""" prefix = tmp_path / "studio_venv" _populate_studio_venv(prefix) out = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix)) @@ -247,22 +236,18 @@ class TestWindowsGpuDetectionAfter5106Fix: site / "nvidia" / "cu13" / "bin" / "x86_64", site / "torch" / "lib", ): - assert ( - str(expected) in out - ), f"resolver missed {expected.relative_to(prefix)}: {out}" + assert str(expected) in out, f"resolver missed {expected.relative_to(prefix)}: {out}" def test_path_assembly_makes_cudart_reachable_without_toolkit(self, tmp_path): """The #5106 scenario: GPU detected, pip nvidia wheels present, - no system CUDA toolkit. cudart must be reachable from PATH, and - from BOTH binary_dir (#5322) and a pip nvidia dir (#5324).""" + no system CUDA toolkit. cudart must be reachable from PATH via + BOTH binary_dir (#5322) and a pip nvidia dir (#5324).""" prefix = tmp_path / "studio_venv" install = tmp_path / "studio_install" _populate_studio_venv(prefix) _populate_studio_install(install, runtime = "13.1") binary_dir = install / "build" / "bin" / "Release" - path_dirs = _build_path_dirs_like_start_llama_server( - binary_dir, prefix, cuda_path = "" - ) + path_dirs = _build_path_dirs_like_start_llama_server(binary_dir, prefix, cuda_path = "") # binary_dir first -- Windows DLL search step 1. assert path_dirs[0] == str( binary_dir @@ -278,15 +263,13 @@ class TestWindowsGpuDetectionAfter5106Fix: ) # Defence in depth: both fix paths contribute cudart. sources = {Path(e).relative_to(tmp_path).parts[0] for e, _ in cudart_locations} - assert ( - "studio_install" in sources - ), f"#5322's cudart drop not reachable: {cudart_locations}" + assert "studio_install" in sources, f"#5322's cudart drop not reachable: {cudart_locations}" assert ( "studio_venv" in sources ), f"#5324's pip nvidia dir not contributing cudart: {cudart_locations}" def test_cublas_and_cublasLt_also_reachable(self, tmp_path): - """ggml-cuda imports cublas64; cublas64 imports cublasLt64. All + """ggml-cuda imports cublas64, which imports cublasLt64. All three must resolve or LoadLibrary returns NULL.""" prefix = tmp_path / "studio_venv" install = tmp_path / "studio_install" @@ -297,12 +280,11 @@ class TestWindowsGpuDetectionAfter5106Fix: for required in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]: reachable = any((Path(d) / required).exists() for d in path_dirs) assert reachable, ( - f"{required} unreachable from PATH; #5106 not fixed.\n" - f"PATH entries: {path_dirs}" + f"{required} unreachable from PATH; #5106 not fixed.\n" f"PATH entries: {path_dirs}" ) def test_no_pip_nvidia_wheels_still_works_via_install_dir(self, tmp_path): - """No pip nvidia wheels (CPU-only torch / unsloth run standalone): + """No pip nvidia wheels (CPU-only torch / standalone unsloth): cudart still resolves via #5322's binary_dir drop.""" prefix = tmp_path / "bare_venv" prefix.mkdir() @@ -310,9 +292,7 @@ class TestWindowsGpuDetectionAfter5106Fix: _populate_studio_install(install, runtime = "13.1") binary_dir = install / "build" / "bin" / "Release" path_dirs = _build_path_dirs_like_start_llama_server(binary_dir, prefix) - assert path_dirs == [ - str(binary_dir) - ], f"bare venv produced unexpected PATH: {path_dirs}" + assert path_dirs == [str(binary_dir)], f"bare venv produced unexpected PATH: {path_dirs}" for required in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]: assert ( binary_dir / required @@ -320,7 +300,7 @@ class TestWindowsGpuDetectionAfter5106Fix: def test_no_install_dir_still_works_via_pip_wheels(self, tmp_path): """Pre-#5322 install (binary_dir lacks cudart): #5324's pip - wheel directories on PATH still resolve cudart.""" + wheel dirs on PATH still resolve cudart.""" prefix = tmp_path / "studio_venv" _populate_studio_venv(prefix) install = tmp_path / "studio_install_pre5322" @@ -336,8 +316,7 @@ class TestWindowsGpuDetectionAfter5106Fix: (rel / fn).write_bytes(b"PE-stub") path_dirs = _build_path_dirs_like_start_llama_server(rel, prefix) cudart_reachable = any( - (Path(d) / "cudart64_12.dll").exists() - or (Path(d) / "cudart64_13.dll").exists() + (Path(d) / "cudart64_12.dll").exists() or (Path(d) / "cudart64_13.dll").exists() for d in path_dirs ) assert cudart_reachable, ( @@ -345,16 +324,15 @@ class TestWindowsGpuDetectionAfter5106Fix: f"on cudart-less install. PATH entries: {path_dirs}" ) cublas_reachable = any( - (Path(d) / "cublas64_12.dll").exists() - or (Path(d) / "cublas64_13.dll").exists() + (Path(d) / "cublas64_12.dll").exists() or (Path(d) / "cublas64_13.dll").exists() for d in path_dirs ) assert cublas_reachable, "cublas unreachable on cudart-less install" def test_pre_pr_scenario_would_have_failed(self, tmp_path): - """Negative control: pre-#5322 + pre-#5324 world leaves cudart + """Negative control: pre-#5322 + pre-#5324 leaves cudart unreachable -- the original failure mode. Confirms the test - actually catches a regression.""" + catches a regression.""" prefix = tmp_path / "studio_venv" _populate_studio_venv(prefix) install = tmp_path / "pre_pr_install" @@ -362,11 +340,10 @@ class TestWindowsGpuDetectionAfter5106Fix: rel.mkdir(parents = True) for fn in ("llama-server.exe", "llama.dll", "ggml-cuda.dll"): (rel / fn).write_bytes(b"PE-stub") - # Pre-PR PATH: binary_dir only. No pip nvidia dirs, no toolkit. + # Pre-PR PATH: binary_dir only, no pip nvidia dirs, no toolkit. pre_pr_path_dirs = [str(rel)] cudart_reachable_pre = any( - (Path(d) / "cudart64_12.dll").exists() - or (Path(d) / "cudart64_13.dll").exists() + (Path(d) / "cudart64_12.dll").exists() or (Path(d) / "cudart64_13.dll").exists() for d in pre_pr_path_dirs ) assert not cudart_reachable_pre, ( @@ -376,9 +353,9 @@ class TestWindowsGpuDetectionAfter5106Fix: class TestWindowsSysPlatformMocked: - """Confirm the win32 branch in start_llama_server is what we test - (not the linux fallback). Patches sys.platform and re-runs the - branch-selecting helper.""" + """Confirm we test the win32 branch in start_llama_server, not the + linux fallback. Patches sys.platform and re-runs the branch-selecting + helper.""" def test_sys_platform_win32_uses_pip_nvidia_resolver(self, monkeypatch, tmp_path): monkeypatch.setattr(sys, "platform", "win32") @@ -387,7 +364,5 @@ class TestWindowsSysPlatformMocked: out = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix)) assert out, f"resolver returned empty under sys.platform=win32: {out}" # cu13 arch dir must be in the output. - cu13_arch = ( - prefix / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" - ) + cu13_arch = prefix / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" assert str(cu13_arch) in out diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py new file mode 100644 index 0000000000..b1c55b61b9 --- /dev/null +++ b/studio/backend/utils/api_errors.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Error-envelope helpers for the OpenAI/Anthropic-compatible ``/v1/*`` API surface. + +FastAPI's defaults emit ``{"detail": ...}`` bodies (status 422 for validation, +``exc.status_code`` for ``HTTPException``). Real OpenAI/Anthropic clients expect +provider-specific error envelopes instead, so this module re-wraps Unsloth's own +client-error responses on the ``/v1/*`` surface: + +- OpenAI surface (``/v1/chat/completions``, ``/v1/completions``, ``/v1/models``, + ``/v1/responses``, ``/v1/embeddings``, ...):: + + {"error": {"message": str, "type": str, "param": None|str, "code": None|str}} + +- Anthropic surface (any path starting with ``/v1/messages``):: + + {"type": "error", "error": {"type": str, "message": str}} + +CRITICAL: the exception handlers installed by :func:`install_api_error_handlers` +are global, but they ONLY transform responses for paths that start with ``/v1/``. +For every other path (``/api/...``, frontend routes) they reproduce FastAPI's +default behavior byte-for-byte, because the Studio frontend depends on the +``{"detail": ...}`` shape for ``/api/*``. + +Public contract (other modules depend on these): + +- ``OPENAI_TYPE_BY_STATUS`` / ``ANTHROPIC_TYPE_BY_STATUS``: status -> type maps. +- ``openai_error_body(message, *, status=400, err_type=None, code=None, param=None)`` +- ``anthropic_error_body(message, *, status=400, err_type=None)`` +- ``is_anthropic_path(path)`` +- ``error_body_for_path(path, message, *, status, err_type=None, code=None, param=None)`` +- ``install_api_error_handlers(app)`` +""" + +from fastapi.encoders import jsonable_encoder +from fastapi.responses import JSONResponse, Response +from fastapi.exceptions import RequestValidationError +from fastapi.utils import is_body_allowed_for_status_code +from starlette.exceptions import HTTPException as StarletteHTTPException + + +# Status-code -> error ``type`` string for the OpenAI error envelope. +OPENAI_TYPE_BY_STATUS = { + 400: "invalid_request_error", + 401: "authentication_error", + 403: "permission_error", + 404: "not_found_error", + 409: "conflict_error", + 413: "invalid_request_error", + 422: "invalid_request_error", + 429: "rate_limit_error", + 500: "api_error", + 502: "api_error", + 503: "api_error", +} + +# Status-code -> error ``type`` string for the Anthropic error envelope. +ANTHROPIC_TYPE_BY_STATUS = { + 400: "invalid_request_error", + 401: "authentication_error", + 403: "permission_error", + 404: "not_found_error", + 409: "conflict_error", + 413: "request_too_large", + 422: "invalid_request_error", + 429: "rate_limit_error", + 500: "api_error", + 502: "api_error", + 503: "api_error", + 529: "overloaded_error", +} + + +def openai_error_body( + message, + *, + status = 400, + err_type = None, + code = None, + param = None, +) -> dict: + """Build an OpenAI-style error envelope. + + Returns ``{"error": {"message", "type", "param", "code"}}``. The ``param`` + and ``code`` keys are always present (value may be ``None``). ``err_type`` + defaults to :data:`OPENAI_TYPE_BY_STATUS` for ``status`` (``"api_error"`` + fallback). + """ + return { + "error": { + "message": str(message), + "type": err_type or OPENAI_TYPE_BY_STATUS.get(status, "api_error"), + "param": param, + "code": code, + } + } + + +def anthropic_error_body( + message, + *, + status = 400, + err_type = None, +) -> dict: + """Build an Anthropic-style error envelope. + + Returns ``{"type": "error", "request_id": None, "error": {"type", "message"}}``. + ``request_id`` is a required (nullable) field on the spec's ErrorResponse; + Studio has no request-id system, so it is null. ``err_type`` defaults to + :data:`ANTHROPIC_TYPE_BY_STATUS` for ``status`` (``"api_error"`` fallback). + """ + return { + "type": "error", + "request_id": None, + "error": { + "type": err_type or ANTHROPIC_TYPE_BY_STATUS.get(status, "api_error"), + "message": str(message), + }, + } + + +def is_anthropic_path(path: str) -> bool: + """True iff ``path`` belongs to the Anthropic surface (``/v1/messages*``).""" + return path.startswith("/v1/messages") + + +def error_body_for_path( + path, + message, + *, + status, + err_type = None, + code = None, + param = None, +) -> dict: + """Dispatch to the correct envelope builder based on ``path``. + + Anthropic surface paths use :func:`anthropic_error_body` (``code``/``param`` + are not part of that envelope and are ignored); all other ``/v1/*`` paths use + :func:`openai_error_body`. + """ + if is_anthropic_path(path): + return anthropic_error_body(message, status = status, err_type = err_type) + return openai_error_body(message, status = status, err_type = err_type, code = code, param = param) + + +def _summarize_validation_errors(errors) -> tuple: + """Derive a readable one-line message and (optional) body param from ``exc.errors()``. + + Returns ``(summary, param)``. ``summary`` is a human-readable string like + ``"messages: Field required"``. ``param`` is the offending body field name when + one can be extracted (used as the OpenAI envelope ``param``), else ``None``. + + Malformed-JSON bodies surface here as ``type == "json_invalid"`` and get a + dedicated message. + """ + if not errors: + return "Invalid request", None + + first = errors[0] + if first.get("type") == "json_invalid": + return "Invalid JSON in request body", None + + loc = first.get("loc", ()) or () + msg = first.get("msg", "Invalid request") + + # Extract the body field name (the loc element after a leading "body"). + param = None + loc_parts = [p for p in loc if p not in ("body",)] + if loc and loc[0] == "body" and loc_parts: + # First non-"body" element that is a field name (string). + for part in loc_parts: + if isinstance(part, str): + param = part + break + + label = ".".join(str(p) for p in loc_parts) if loc_parts else ".".join(str(p) for p in loc) + summary = f"{label}: {msg}" if label else str(msg) + return summary, param + + +def install_api_error_handlers(app) -> None: + """Register validation + HTTPException handlers that emit ``/v1/*`` envelopes. + + Both handlers are global but only transform responses for paths starting with + ``/v1/``. Non-``/v1/`` paths reproduce FastAPI's default ``{"detail": ...}`` + behavior exactly so the Studio frontend keeps working. + """ + + @app.exception_handler(RequestValidationError) + async def _handle_validation_error(request, exc): + path = request.url.path + if path.startswith("/v1/"): + summary, param = _summarize_validation_errors(exc.errors()) + return JSONResponse( + status_code = 400, + content = error_body_for_path(path, summary, status = 400, param = param), + ) + # Default FastAPI behavior for every other path. + return JSONResponse( + status_code = 422, + content = {"detail": jsonable_encoder(exc.errors())}, + ) + + @app.exception_handler(StarletteHTTPException) + async def _handle_http_exception(request, exc): + path = request.url.path + headers = getattr(exc, "headers", None) + # Statuses like 204/304/1xx must not carry a body — mirror FastAPI's + # default http_exception_handler, which returns a bodiless Response. + if not is_body_allowed_for_status_code(exc.status_code): + return Response(status_code = exc.status_code, headers = headers) + if path.startswith("/v1/"): + detail = exc.detail + # Already a fully-formed envelope: pass through untouched. + if isinstance(detail, dict) and ("error" in detail or detail.get("type") == "error"): + return JSONResponse( + status_code = exc.status_code, + content = detail, + headers = headers, + ) + # A dict carrying our individual fields. + if isinstance(detail, dict): + message = detail.get("message", detail) + err_type = detail.get("type") + code = detail.get("code") + param = detail.get("param") + else: + # Plain message string (the common HTTPException case). + message = detail + err_type = None + code = None + param = None + return JSONResponse( + status_code = exc.status_code, + content = error_body_for_path( + path, + message, + status = exc.status_code, + err_type = err_type, + code = code, + param = param, + ), + headers = headers, + ) + # Default FastAPI behavior for every other path. + return JSONResponse( + status_code = exc.status_code, + content = {"detail": exc.detail}, + headers = headers, + ) diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py index 4c8e6239a0..210735973d 100644 --- a/studio/backend/utils/cache_cleanup.py +++ b/studio/backend/utils/cache_cleanup.py @@ -1,14 +1,12 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Utility for cleaning up the Unsloth compiled cache directory. +"""Clean up the Unsloth compiled cache directory. -The unsloth_compiled_cache is created by unsloth_zoo/compiler.py during -FastModel.from_pretrained() and contains model-type-specific compiled Python -files. It should be selectively cleared between model loads to avoid stale -artefacts, while preserving model-agnostic components (like Trainers) needed -by spawned subprocesses. +unsloth_compiled_cache (created by unsloth_zoo/compiler.py during +FastModel.from_pretrained) holds model-type-specific compiled files. Clear it +selectively between model loads, preserving model-agnostic components (Trainers) +that spawned subprocesses need. """ import shutil @@ -38,9 +36,8 @@ def get_existing_cache_dirs() -> List[Path]: def register_compiled_cache_on_path() -> None: """Add all existing compiled-cache directories to sys.path and PYTHONPATH. - This ensures spawned workers (on platforms using the 'spawn' start method, - i.e. Windows and macOS) can import dynamically compiled modules such as - UnslothSFTTrainer. + Ensures spawned workers (on 'spawn'-start platforms, i.e. Windows and macOS) + can import dynamically compiled modules such as UnslothSFTTrainer. """ import os import sys @@ -48,8 +45,8 @@ def register_compiled_cache_on_path() -> None: pypath = os.environ.get("PYTHONPATH", "") pypath_entries = [p for p in pypath.split(os.pathsep) if p] - # Iterate in reverse so that earlier _CACHE_DIRS entries (higher priority) - # are inserted last and therefore end up first in sys.path / PYTHONPATH. + # Iterate in reverse so earlier _CACHE_DIRS entries (higher priority) are + # inserted last and thus end up first in sys.path / PYTHONPATH. for cache_dir in reversed(get_existing_cache_dirs()): resolved = str(cache_dir.resolve()) if resolved not in sys.path: @@ -65,7 +62,7 @@ def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None) Remove compiled files from the cache directory (idempotent). Args: - preserve_patterns: A list of glob patterns for files to keep + preserve_patterns: glob patterns for files to keep (e.g., ["Unsloth*Trainer.py"]). If None or empty, the entire cache directory is deleted (legacy behavior). """ @@ -75,13 +72,11 @@ def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None) if preserve_patterns: logger.info( - f"Cleaning unsloth compiled cache (preserving {preserve_patterns}): " - f"{cache_dir}" + f"Cleaning unsloth compiled cache (preserving {preserve_patterns}): " f"{cache_dir}" ) for item in cache_dir.iterdir(): if item.is_file(): - # Check if the file matches any of the patterns we want to keep preserve = any(item.match(pattern) for pattern in preserve_patterns) if not preserve: try: @@ -93,6 +88,6 @@ def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None) # Always clear __pycache__ and other subdirectories shutil.rmtree(item, ignore_errors = True) else: - # Legacy behavior: nuke the entire directory + # Legacy: remove the entire directory logger.info(f"Removing unsloth compiled cache: {cache_dir}") shutil.rmtree(cache_dir, ignore_errors = True) diff --git a/studio/backend/utils/cpu_threads.py b/studio/backend/utils/cpu_threads.py new file mode 100644 index 0000000000..4ed0021054 --- /dev/null +++ b/studio/backend/utils/cpu_threads.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Early CPU thread-pool configuration for Studio processes.""" + +import os +from typing import MutableMapping, Optional + + +_THREAD_POOL_ENV_VARS = ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "NUMEXPR_NUM_THREADS", +) + + +def configure_cpu_threads(env: Optional[MutableMapping[str, str]] = None) -> None: + """Apply ``UNSLOTH_CPU_THREADS`` to native CPU pools when configured. + + Must run before importing libraries that initialize an OpenMP or BLAS + pool. Library-specific vars are left untouched so users can override a + single runtime independently. + """ + environ = os.environ if env is None else env + configured = environ.get("UNSLOTH_CPU_THREADS", "").strip() + if not configured: + return + + try: + thread_count = int(configured) + except ValueError as exc: + raise ValueError("UNSLOTH_CPU_THREADS must be a positive integer") from exc + if thread_count < 1: + raise ValueError("UNSLOTH_CPU_THREADS must be a positive integer") + + value = str(thread_count) + for variable in _THREAD_POOL_ENV_VARS: + environ.setdefault(variable, value) diff --git a/studio/backend/utils/datasets/__init__.py b/studio/backend/utils/datasets/__init__.py index 7988b09972..caa471bde5 100644 --- a/studio/backend/utils/datasets/__init__.py +++ b/studio/backend/utils/datasets/__init__.py @@ -1,22 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Dataset utilities package. +"""Dataset utilities for LLM/VLM fine-tuning: detection, conversion, templating, collators, mappings.""" -This package provides utilities for dataset format detection, conversion, -and processing for LLM and VLM fine-tuning workflows. - -Modules: -- format_detection: Detect dataset formats (Alpaca, ShareGPT, ChatML) -- format_conversion: Convert between dataset formats -- chat_templates: Apply chat templates to datasets -- vlm_processing: Vision-Language Model processing utilities -- data_collators: Custom data collators for training -- model_mappings: Model-to-template mapping constants -""" - -# Format detection from .format_detection import ( detect_dataset_format, detect_custom_format_heuristic, @@ -24,7 +10,6 @@ from .format_detection import ( detect_vlm_dataset_structure, ) -# Format conversion from .format_conversion import ( standardize_chat_format, convert_chatml_to_alpaca, @@ -34,7 +19,6 @@ from .format_conversion import ( convert_sharegpt_with_images_to_vlm_format, ) -# Chat templates from .chat_templates import ( apply_chat_template_to_dataset, get_dataset_info_summary, @@ -42,19 +26,16 @@ from .chat_templates import ( DEFAULT_ALPACA_TEMPLATE, ) -# VLM processing from .vlm_processing import ( generate_smart_vlm_instruction, ) -# Data collators from .data_collators import ( DataCollatorSpeechSeq2SeqWithPadding, DeepSeekOCRDataCollator, VLMDataCollator, ) -# Model mappings (constants) from .model_mappings import ( TEMPLATE_TO_MODEL_MAPPER, MODEL_TO_TEMPLATE_MAPPER, @@ -62,15 +43,13 @@ from .model_mappings import ( is_gpt_oss_model_name, ) -# Legacy imports from the original dataset_utils.py for backward compatibility -# These functions have not yet been refactored into separate modules +# Legacy dataset_utils.py imports kept for backward compat from .dataset_utils import ( check_dataset_format, format_and_template_dataset, format_dataset, ) -# Public API __all__ = [ # Detection "detect_dataset_format", diff --git a/studio/backend/utils/datasets/cache_safe.py b/studio/backend/utils/datasets/cache_safe.py new file mode 100644 index 0000000000..e629210f33 --- /dev/null +++ b/studio/backend/utils/datasets/cache_safe.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Permission-safe wrapper around datasets.load_dataset. + +A shared HF datasets cache can contain subtrees owned by another user (for +example populated by an earlier root-run job). datasets then raises +"[Errno 13] Permission denied: ..._builder.lock" while locking the cached +builder, killing the training run even though the dataset itself is fine. +Retry such loads in a Studio-owned cache so the run proceeds; the worst case +is one rebuild of the dataset in the fallback location. +""" + +import logging +import os + +from utils.paths.storage_roots import cache_root + +logger = logging.getLogger(__name__) + + +def studio_datasets_cache() -> str: + path = cache_root() / "hf-datasets" + path.mkdir(parents = True, exist_ok = True) + return str(path) + + +def load_dataset_cache_safe(*args, **kwargs): + """datasets.load_dataset, retried in a Studio-owned cache on EACCES.""" + from datasets import load_dataset + try: + return load_dataset(*args, **kwargs) + except PermissionError as error: + fallback = studio_datasets_cache() + logger.warning( + "HF datasets cache is not writable (%s); rebuilding in %s", + error, + fallback, + ) + kwargs["cache_dir"] = fallback + # Nested builders consult the env var while the load runs; restore it + # after so other datasets keep trying the shared cache first. + old_env = os.environ.get("HF_DATASETS_CACHE") + os.environ["HF_DATASETS_CACHE"] = fallback + try: + return load_dataset(*args, **kwargs) + finally: + if old_env is None: + os.environ.pop("HF_DATASETS_CACHE", None) + else: + os.environ["HF_DATASETS_CACHE"] = old_env diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index cfdd811853..82a30fd55b 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Chat template application utilities for dataset processing. +"""Chat template utilities for dataset processing. -This module contains functions for applying chat templates to datasets -and generating dataset info summaries. +Apply chat templates to datasets and generate dataset info summaries. """ from .format_detection import detect_dataset_format, detect_multimodal_dataset, detect_custom_format_heuristic @@ -46,30 +44,25 @@ def _chat_template_kwargs() -> dict: def get_tokenizer_chat_template(tokenizer, model_name): - """ - Gets appropriate chat template for tokenizer based on model. - Uses Unsloth's get_chat_template if model is in the mapper. + """Apply a chat template to the tokenizer, using Unsloth's + get_chat_template when the model class name is in the mapper. Args: tokenizer: HuggingFace tokenizer model_name: Model class name (e.g., "Gemma3ForCausalLM") Returns: - tokenizer: Tokenizer with appropriate chat template applied + tokenizer with the chat template applied """ try: from unsloth.chat_templates import get_chat_template except ImportError: - # Unsloth not available, return tokenizer as-is return tokenizer - # Normalize model_name to lowercase for matching model_name_lower = model_name.lower() - # Check if model matches any template in mapper matched_template = None - # Direct match in MODEL_TO_TEMPLATE_MAPPER if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: matched_template = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] logger.info(f"📝 Applying Unsloth chat template: {matched_template}") @@ -83,7 +76,6 @@ def get_tokenizer_chat_template(tokenizer, model_name): logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}") logger.info(f" Falling back to tokenizer's default chat template") else: - # Check if tokenizer actually has a chat_template set has_chat_template = ( hasattr(tokenizer, 'chat_template') and tokenizer.chat_template is not None @@ -91,7 +83,7 @@ def get_tokenizer_chat_template(tokenizer, model_name): if has_chat_template: logger.info(f"📝 Using tokenizer's own chat template (no Unsloth template match)") else: - # Base model with no chat template — apply default ChatML + # Base model with no chat template: apply default ChatML. logger.info(f"📝 No chat template found — applying default ChatML template (base model)") try: tokenizer = get_chat_template( @@ -107,9 +99,7 @@ def get_tokenizer_chat_template(tokenizer, model_name): def get_dataset_info_summary(dataset_info): - """ - Returns a human-readable summary for UI display. - """ + """Return a human-readable summary for UI display.""" detected_format = dataset_info["detected_format"] final_format = dataset_info["final_format"] @@ -146,15 +136,14 @@ def apply_chat_template_to_dataset( num_proc = None, progress_callback = None, ): - """ - Applies chat template to dataset based on its format. + """Apply the chat template to a dataset based on its format. Args: dataset_info: Output from format_dataset() with metadata tokenizer: Tokenizer with chat template custom_prompt_template: Optional string template for custom formatting - add_eos_token: If True, appends tokenizer.eos_token to each text - remove_bos_prefix: If True, removes '' prefix (for Gemma, etc.) + add_eos_token: If True, append tokenizer.eos_token to each text + remove_bos_prefix: If True, remove '' prefix (Gemma, etc.) custom_format_mapping: Dict mapping custom columns to standard format batch_size: Batch size for processing num_proc: Number of processes @@ -170,7 +159,6 @@ def apply_chat_template_to_dataset( warnings = list(dataset_info.get("warnings", [])) errors = [] - # Get EOS token if needed eos_token = "" if add_eos_token: if hasattr(tokenizer, 'eos_token') and tokenizer.eos_token: @@ -180,9 +168,8 @@ def apply_chat_template_to_dataset( # CUSTOM FORMAT MAPPING (for non-standard datasets) if final_format == "unknown": - # Try auto-detection if no custom mapping provided if custom_format_mapping is None and auto_detect_mapping: - # Check if format_dataset already tried and failed + # Skip if format_dataset already tried and failed. if not dataset_info.get("auto_detection_attempted", False): custom_format_mapping = detect_custom_format_heuristic(dataset) if custom_format_mapping: @@ -196,7 +183,7 @@ def apply_chat_template_to_dataset( "errors": errors } else: - # Already failed once in format_dataset, don't retry + # Already failed once in format_dataset; don't retry. errors.append( "Format remains unknown after detection attempts. " "Please provide custom_format_mapping to specify column roles manually." @@ -216,7 +203,7 @@ def apply_chat_template_to_dataset( conversations = [] num_examples = len(examples[list(examples.keys())[0]]) - # Only preserve unmapped columns if auto-detected + # Preserve unmapped columns only if auto-detected. preserved_columns = {} if not is_user_provided: all_columns = set(examples.keys()) @@ -236,10 +223,10 @@ def apply_chat_template_to_dataset( content = examples[col_name][i] if is_user_provided: - # User explicitly mapped - include even if empty + # User-mapped: include even if empty. convo.append({"role": role, "content": str(content) if content else ""}) else: - # Auto-detected - skip empty + # Auto-detected: skip empty. if content and str(content).strip(): convo.append({"role": role, "content": str(content)}) @@ -252,7 +239,6 @@ def apply_chat_template_to_dataset( try: dataset = dataset.map(_apply_custom_mapping, batched = True, batch_size = batch_size) - # Update to use conversations format final_format = "chatml_conversations" chat_column = "conversations" is_standardized = True @@ -269,8 +255,7 @@ def apply_chat_template_to_dataset( # ALPACA FORMAT if final_format == "alpaca": - # Set alpaca chat template on tokenizer for saving (if not already set) - # This ensures the template is saved with the model for inference + # Set alpaca chat template (if unset) so it's saved for inference. if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template): try: from unsloth.chat_templates import get_chat_template @@ -283,7 +268,6 @@ def apply_chat_template_to_dataset( except Exception as e: logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}") - # Use custom template if provided def _format_alpaca_custom(examples): texts = [] for i in range(len(examples["instruction"])): @@ -349,7 +333,7 @@ def apply_chat_template_to_dataset( if not is_standardized: warnings.append("Dataset may not be fully standardized") - # Apply Unsloth chat template if model matches + # Apply Unsloth chat template if the model matches. if model_name: tokenizer = get_tokenizer_chat_template(tokenizer, model_name) @@ -398,7 +382,7 @@ def apply_chat_template_to_dataset( dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = f"Applying chat template to {final_format}" - # Monitor tqdm progress from dataset.map() and relay to callback + # Monitor dataset.map() tqdm progress and relay it. _tqdm_monitor_stop = None if progress_callback and not _is_torch_iterable: import threading diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py index 687da74c21..9bfb60ba17 100644 --- a/studio/backend/utils/datasets/data_collators.py +++ b/studio/backend/utils/datasets/data_collators.py @@ -1,12 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Data collators for dataset processing. - -This module contains custom data collators for training, -particularly for VLM/OCR processing. -""" +"""Custom training data collators, particularly for VLM/OCR processing.""" from dataclasses import dataclass from typing import Any, List, Optional, Union @@ -20,27 +15,21 @@ class DataCollatorSpeechSeq2SeqWithPadding: """ Data collator for Whisper speech-to-text training. - Pads input features (audio) and label sequences (text) separately, - masks padding in labels with -100, and strips leading BOS token. - Mirrors the collator from the Whisper.ipynb notebook. + Pads audio input features and text labels separately, masks label padding + with -100, and strips the leading BOS token. Mirrors the Whisper.ipynb + notebook collator. """ processor: Any def __call__(self, features: List[dict]) -> dict: - input_features = [ - {"input_features": feature["input_features"]} for feature in features - ] - batch = self.processor.feature_extractor.pad( - input_features, return_tensors = "pt" - ) + input_features = [{"input_features": feature["input_features"]} for feature in features] + batch = self.processor.feature_extractor.pad(input_features, return_tensors = "pt") label_features = [{"input_ids": feature["labels"]} for feature in features] labels_batch = self.processor.tokenizer.pad(label_features, return_tensors = "pt") - labels = labels_batch["input_ids"].masked_fill( - labels_batch.attention_mask.ne(1), -100 - ) + labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100) if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item(): labels = labels[:, 1:] @@ -51,13 +40,10 @@ class DataCollatorSpeechSeq2SeqWithPadding: @dataclass class DeepSeekOCRDataCollator: - """ - Data collator for DeepSeek OCR VLM training. + """Data collator for DeepSeek OCR VLM training. - Handles: - - Image processing via processor - - Text tokenization - - Proper label masking for instruction fine-tuning + Handles image processing, text tokenization, and label masking for + instruction fine-tuning. """ processor: Any # Qwen2VLProcessor or similar @@ -77,7 +63,6 @@ class DeepSeekOCRDataCollator: """ from PIL import Image - # Extract messages and images all_messages = [] all_images = [] @@ -85,7 +70,6 @@ class DeepSeekOCRDataCollator: messages = sample["messages"] all_messages.append(messages) - # Extract PIL images from content for msg in messages: content = msg.get("content", []) if isinstance(content, list): @@ -95,9 +79,7 @@ class DeepSeekOCRDataCollator: if img is not None and hasattr(img, "size"): # PIL Image all_images.append(img) - # Process with the VL processor try: - # Qwen2VL style processing texts = [ self.processor.apply_chat_template( msgs, tokenize = False, add_generation_prompt = False @@ -105,7 +87,6 @@ class DeepSeekOCRDataCollator: for msgs in all_messages ] - # Process with images inputs = self.processor( text = texts, images = all_images if all_images else None, @@ -115,10 +96,7 @@ class DeepSeekOCRDataCollator: max_length = self.max_length, ) - # Create labels (mask input, keep output) labels = inputs["input_ids"].clone() - - # Simple masking: mask padding tokens labels[labels == self.processor.tokenizer.pad_token_id] = self.ignore_index inputs["labels"] = labels @@ -132,24 +110,15 @@ class DeepSeekOCRDataCollator: @dataclass class VLMDataCollator: - """ - Generic VLM data collator that works with various processors. - - Supports: - - Qwen2VL - - LLaVA - - Other VL models with compatible processors - """ + """Generic VLM data collator for various processors (Qwen2VL, LLaVA, etc.).""" processor: Any max_length: int = 2048 ignore_index: int = -100 - mask_input_tokens: bool = True # Whether to mask user tokens in labels + mask_input_tokens: bool = True # Mask user tokens in labels def __call__(self, batch: List[dict]) -> dict: - """ - Collate a batch of VLM samples. - """ + """Collate a batch of VLM samples.""" all_messages = [] all_images = [] @@ -157,7 +126,6 @@ class VLMDataCollator: messages = sample.get("messages", []) all_messages.append(messages) - # Extract images for msg in messages: content = msg.get("content", []) if isinstance(content, list): @@ -167,15 +135,11 @@ class VLMDataCollator: if img is not None: all_images.append(img) - # Apply chat template texts = [ - self.processor.apply_chat_template( - msgs, tokenize = False, add_generation_prompt = False - ) + self.processor.apply_chat_template(msgs, tokenize = False, add_generation_prompt = False) for msgs in all_messages ] - # Process inputs inputs = self.processor( text = texts, images = all_images if all_images else None, @@ -185,10 +149,9 @@ class VLMDataCollator: max_length = self.max_length, ) - # Create labels labels = inputs["input_ids"].clone() - # Mask padding + # Mask padding. if hasattr(self.processor, "tokenizer"): pad_token_id = self.processor.tokenizer.pad_token_id else: diff --git a/studio/backend/utils/datasets/dataset_none_detect.py b/studio/backend/utils/datasets/dataset_none_detect.py new file mode 100644 index 0000000000..a2fd8ef667 --- /dev/null +++ b/studio/backend/utils/datasets/dataset_none_detect.py @@ -0,0 +1,855 @@ +""" +Detect None/empty content turns in conversation datasets. Reports findings +without modifying data. + +Usage: + from .dataset_none_detect import scan_dataset, print_report + stats = scan_dataset(dataset) # auto-detect + scan + stats = scan_dataset(dataset, fmt="chatml") # explicit format + print_report(stats, stats["format"]) + +Dependencies: only `datasets` (already in studio/unsloth) + stdlib. + +Supported formats (via FORMAT_REGISTRY): + alpaca instruction/output instruction + output must be set + chatml messages/conversations/texts role + content per turn + sharegpt conversations from/value per turn + gptoss messages (alias: gpt-oss) role/content; has a developer turn + +Any role/content chat template matches chatml, so new templates need no change; +add a FORMAT_REGISTRY entry only for a genuinely new column/turn shape. +""" + +from datasets import Dataset + +# --------------------------------------------------------------------------- +# Conversation column probing (shared by detection + scanning) +# --------------------------------------------------------------------------- + +# Candidate column names for conversational datasets, checked in priority order. +CONVERSATION_COLUMNS = ("messages", "conversations", "texts") + +# Minimum turn key sets identifying a column as conversational (not e.g. messages=[{"id":1}]). +_CHAT_KEY_SETS = (frozenset({"role", "content"}), frozenset({"from", "value"})) + + +def _probe_conversation(dataset: Dataset, candidates = None): + """ + Probe a dataset for its conversation column and turn structure. + + candidates - column names to try, in priority order. + Defaults to CONVERSATION_COLUMNS when None. + + Returns a dict with: + column - conversation column found + turn_keys - keys present in the first turn dict + roles - all role values seen across the first few samples + + Returns None if no conversation column is found. + """ + if candidates is None: + candidates = CONVERSATION_COLUMNS + columns = set(dataset.column_names) + # Remember the first all-corrupt candidate, but keep probing: a later column + # may be healthy and win (e.g. bad messages, good conversations). + all_corrupt_fallback = None + for col in candidates: + if col not in columns: + continue + # Scan up to 100 rows - row 0 alone may be empty/malformed. + first = None + for i in range(min(len(dataset), 100)): + sample = dataset[i][col] + if not isinstance(sample, list) or len(sample) == 0: + continue + # Skip non-dict leading turns (e.g. [None, {"role": ...}]). + first_turn = next((t for t in sample if isinstance(t, dict)), None) + if first_turn is not None: + first = first_turn + break + if first is None: + # No usable dict turn in 100 rows. Record an all_corrupt fallback, + # plausible only with turn-shaped data (None cell or list of dict/None + # turns); a later plausible candidate upgrades a non-plausible one. + if all_corrupt_fallback is None or not all_corrupt_fallback.get("has_plausible_turns"): + has_plausible_turns = False + for i in range(min(len(dataset), 100)): + cell = dataset[i][col] + if cell is None: + has_plausible_turns = True + break + # A struct-typed cell (single dict, not a list) is metadata, + # not chat: leave it for "unknown format", matching + # format_detection.py. + if isinstance(cell, list): + # Plausible only if the list holds a dict/None turn; + # empty lists and list-of-strings are not chat data. + if any(t is None or isinstance(t, dict) for t in cell): + has_plausible_turns = True + break + all_corrupt_fallback = { + "column": col, + "turn_keys": set(), + "roles": set(), + "all_corrupt": True, + "has_plausible_turns": has_plausible_turns, + } + continue + + # Use the same 100-row window to gather keys/roles. + turn_keys = set() + roles = set() + for i in range(min(len(dataset), 100)): + conv = dataset[i][col] + if isinstance(conv, list): + for t in conv: + if isinstance(t, dict): + turn_keys.update(t.keys()) + r = t.get("role") or t.get("from") + if r: + roles.add(str(r)) + # Column lacks a full chat key pair. If it has a conversational key + # (role/from/content/value) it is a corrupt-but-real chat column, so + # save a plausible fallback for find_none_chatml to flag. Pure metadata + # (e.g. [{"id":1}]) is not plausible, so a later real-but-corrupt column + # (e.g. conversations=None) can still win. + _CONV_KEYS = {"role", "from", "content", "value"} + if not any(keys <= turn_keys for keys in _CHAT_KEY_SETS): + schema_less_plausible = bool(turn_keys & _CONV_KEYS) + if all_corrupt_fallback is None or not all_corrupt_fallback.get("has_plausible_turns"): + all_corrupt_fallback = { + "column": col, + "turn_keys": turn_keys, + "roles": roles, + "all_corrupt": True, + "has_plausible_turns": schema_less_plausible, + } + continue + return {"column": col, "turn_keys": turn_keys, "roles": roles} + # No healthy column found; return the all_corrupt fallback if any. + return all_corrupt_fallback + + +# None-detection helpers +# --------------------------------------------------------------------------- + + +def is_none_or_empty(value) -> bool: + """True if value is None, empty string, whitespace-only, or an empty/whitespace-only VLM content block list.""" + if value is None: + return True + if isinstance(value, str): + # Treat zero-width/BOM chars (U+FEFF/200B/200C/200D/2060) as empty; + # they render invisibly. Two-pass strip (ws, invisibles, ws) catches + # mixed cases like "\u200b \u200b". + stripped = value.strip().strip("\ufeff\u200b\u200c\u200d\u2060").strip() + if not stripped: + return True + if isinstance(value, list): + # VLM content blocks, e.g. [{"type":"text",...}, {"type":"image",...}]. + # Empty list -> empty. A non-text block (image/audio/tool) is real + # content; only flag when every text block is blank and no such block + # exists (an image-only turn is valid). + if len(value) == 0: + return True + # No dict blocks (e.g. [None], [' ']) -> malformed/empty. + dict_blocks = [item for item in value if isinstance(item, dict)] + if not dict_blocks: + return True + non_text_blocks = [item for item in dict_blocks if item.get("type") != "text"] + if non_text_blocks: + return False + text_values = [item.get("text") for item in dict_blocks if item.get("type") == "text"] + if text_values and all( + t is None + or ( + isinstance(t, str) and not t.strip().strip("\ufeff\u200b\u200c\u200d\u2060").strip() + ) + for t in text_values + ): + return True + return False + + +def _classify_empty(value) -> str: + """Return a human-readable label for why this value is considered empty.""" + if value is None: + return "None" + if isinstance(value, str): + if len(value) == 0: + return "empty_string" + # Only-whitespace or only-invisible (BOM/zero-width) strings render empty. + if not value.strip().strip("\ufeff\u200b\u200c\u200d\u2060").strip(): + return "whitespace_only" + if isinstance(value, list): + # Mirrors the VLM/OpenAI content-block handling in is_none_or_empty. + if len(value) == 0: + return "empty_list" + return "empty_vlm_content" + return "valid" # unreachable if is_none_or_empty was True + + +# --------------------------------------------------------------------------- +# Alpaca detection +# --------------------------------------------------------------------------- + + +def find_none_alpaca(dataset: Dataset) -> dict: + """ + Scan alpaca dataset for None/empty instruction or output fields. + Returns a stats dict with a detailed 'findings' list. + """ + stats = { + "total_rows": len(dataset), + "none_instruction": 0, + "none_output": 0, + "bad_row_indices": [], + "findings": [], # [{row, field, value_type, raw_value}, ...] + } + + for i, row in enumerate(dataset): + bad = False + for field in ("instruction", "output"): + val = row.get(field) + if is_none_or_empty(val): + stats[f"none_{field}"] = stats.get(f"none_{field}", 0) + 1 + bad = True + stats["findings"].append( + { + "row_index": i, + "field": field, + "value_type": _classify_empty(val), + "raw_value": repr(val), + } + ) + if bad: + stats["bad_row_indices"].append(i) + + return stats + + +# --------------------------------------------------------------------------- +# ChatML / conversational detection +# --------------------------------------------------------------------------- + + +def find_none_chatml(dataset: Dataset, col: str = None) -> dict: + """ + Scan chatml/sharegpt/gptoss dataset for turns with None/empty content. + Auto-detects the conversation column if col=None. + + Returns a stats dict with a complete 'findings' list - one entry per bad + turn with row_index, turn_index, role, value_type, and raw_value. + """ + if col is None: + # Reuse _probe_conversation so the all_corrupt path is handled here too. + _cinfo = _probe_conversation(dataset) + if _cinfo is not None: + col = _cinfo["column"] + + if col is None or col not in dataset.column_names: + raise ValueError( + f"No conversation column found. " + f"Expected one of {CONVERSATION_COLUMNS}, got columns: {dataset.column_names}" + ) + + stats = { + "total_rows": len(dataset), + "column": col, + "rows_with_none_turns": 0, + "total_none_turns": 0, + "none_by_role": {}, # role -> count of None turns + "none_by_type": {}, # "None" | "empty_string" | "whitespace_only" -> count + "rows_all_none": 0, # rows where every turn is bad + "bad_row_indices": [], # every row index that has at least one bad turn + "findings": [], # detailed per-turn list + } + + for i, row in enumerate(dataset): + conversation = row[col] + if not isinstance(conversation, list): + # Non-list conversation: unusable for training, flag as bad. + vtype = "None" if conversation is None else "invalid_type" + stats["bad_row_indices"].append(i) + stats["rows_with_none_turns"] += 1 + stats["total_none_turns"] += 1 + stats["rows_all_none"] += 1 + stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1 + stats["none_by_type"][vtype] = stats["none_by_type"].get(vtype, 0) + 1 + stats["findings"].append( + { + "row_index": i, + "turn_index": 0, + "role": "unknown", + "value_type": vtype, + "raw_value": repr(conversation), + } + ) + continue + + if len(conversation) == 0: + # Zero-turn conversation: flag so it doesn't scan as clean. + stats["bad_row_indices"].append(i) + stats["rows_with_none_turns"] += 1 + stats["total_none_turns"] += 1 + stats["rows_all_none"] += 1 + stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1 + stats["none_by_type"]["empty_conversation"] = ( + stats["none_by_type"].get("empty_conversation", 0) + 1 + ) + stats["findings"].append( + { + "row_index": i, + "turn_index": 0, + "role": "unknown", + "value_type": "empty_conversation", + "raw_value": "[]", + } + ) + continue + + row_findings = [] + for turn_idx, turn in enumerate(conversation): + # Non-dict turn - record it rather than crash or silently skip. + if not isinstance(turn, dict): + row_findings.append( + { + "row_index": i, + "turn_index": turn_idx, + "role": "unknown", + "value_type": "None" if turn is None else "invalid_type", + "raw_value": repr(turn), + } + ) + stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1 + vtype = "None" if turn is None else "invalid_type" + stats["none_by_type"][vtype] = stats["none_by_type"].get(vtype, 0) + 1 + continue + # Explicit None check so falsy roles (0, "", False) are kept, not + # collapsed to "unknown". + r = turn.get("role") + if r is None: + r = turn.get("from") + if r is None: + role = "unknown" + elif isinstance(r, str): + role = r + else: + role = str(r) + # Pick the content key: from+value -> value (ShareGPT, even if role + # is set); role -> content (or value); from only -> value (None when + # missing, so it is flagged); neither -> content then value. + if "from" in turn and "value" in turn: + content = turn.get("value") + elif "role" in turn: + content = turn.get("content") if "content" in turn else turn.get("value") + elif "from" in turn: + content = turn.get("value") + else: + content = turn.get("content") if "content" in turn else turn.get("value") + # Assistant tool-call turns carry empty content + tool_calls and are + # valid; the exemption is assistant-only. + if is_none_or_empty(content) and not (role == "assistant" and turn.get("tool_calls")): + vtype = _classify_empty(content) + row_findings.append( + { + "row_index": i, + "turn_index": turn_idx, + "role": role, + "value_type": vtype, + "raw_value": repr(content), + } + ) + stats["none_by_role"][role] = stats["none_by_role"].get(role, 0) + 1 + stats["none_by_type"][vtype] = stats["none_by_type"].get(vtype, 0) + 1 + + if row_findings: + stats["rows_with_none_turns"] += 1 + stats["total_none_turns"] += len(row_findings) + stats["bad_row_indices"].append(i) + stats["findings"].extend(row_findings) + + if len(row_findings) == len(conversation): + stats["rows_all_none"] += 1 + + return stats + + +# --------------------------------------------------------------------------- +# Convenience wrappers per format (all delegate to the same scan logic) +# --------------------------------------------------------------------------- + + +def find_none_sharegpt(dataset: Dataset, col: str = None) -> dict: + """ShareGPT uses 'from'/'value' keys - same scan logic handles both.""" + if col is None: + # ShareGPT lives in 'conversations'; probe only that column so a corrupt + # one is still scanned, not replaced by healthy 'messages' (P1 fix). + conv_info = _probe_conversation(dataset, candidates = ("conversations",)) + if conv_info is None: + raise ValueError( + f"No valid conversation column found in {dataset.column_names}. " + "Expected a 'conversations' column with 'from'/'value' or 'role'/'content' turn keys." + ) + col = conv_info["column"] + return find_none_chatml(dataset, col = col) + + +def find_none_gptoss(dataset: Dataset, col: str = None) -> dict: + """gptoss: role/content plus optional thinking/tool_calls. Only content checked.""" + if col is None: + # gptoss lives in 'messages': target it whenever present (even if + # corrupt); fall back to 'conversations' only if 'messages' is absent. + if "messages" in dataset.column_names: + conv_info = _probe_conversation(dataset, candidates = ("messages",)) + else: + conv_info = _probe_conversation(dataset, candidates = ("conversations",)) + if conv_info is None: + raise ValueError( + f"No valid conversation column found in {dataset.column_names}. " + "Expected a 'messages' or 'conversations' column with 'role'/'content' turn keys." + ) + col = conv_info["column"] + return find_none_chatml(dataset, col = col) + + +# --------------------------------------------------------------------------- +# Format registry - first match wins; detect_format() auto-scales. +# Each entry: name (label/--format value), match(dataset, conv_info) -> bool, +# scan (find_none_* function). Put specific formats before general ones +# (gptoss before chatml, since gptoss is chatml with a 'developer' role). +# To add a format: write find_none_() (or reuse find_none_chatml) and +# append an entry; detect_format(), --format, and scan_dataset() pick it up. +# --------------------------------------------------------------------------- + +FORMAT_REGISTRY = [ + { + "name": "alpaca", + # instruction/output present and no usable chat column: either none + # exists, or the only one is fully corrupt (e.g. a stray all-None or + # metadata `messages` column). A healthy chat column falls through to + # the conversational scanners below. + "match": lambda ds, conv: ( + {"instruction", "output"}.issubset(ds.column_names) + and (conv is None or conv.get("all_corrupt")) + ), + "scan": find_none_alpaca, + }, + { + "name": "gptoss", + "match": lambda ds, conv: ( + conv is not None + and {"role", "content"} <= conv["turn_keys"] + and "developer" in conv["roles"] + ), + "scan": find_none_gptoss, + }, + { + "name": "sharegpt", + "match": lambda ds, conv: (conv is not None and {"from", "value"} <= conv["turn_keys"]), + "scan": find_none_sharegpt, + }, + { + "name": "chatml", + "match": lambda ds, conv: ( + conv is not None + and ( + {"role", "content"} <= conv["turn_keys"] + # all_corrupt: column found but every row malformed; require + # has_plausible_turns so scalar/string columns aren't chatml. + or (conv.get("all_corrupt") and conv.get("has_plausible_turns")) + ) + ), + "scan": find_none_chatml, + }, +] + +# Derived list of known format names (used by CLI --format choices). +FORMAT_NAMES = [entry["name"] for entry in FORMAT_REGISTRY] + +# Documented aliases accepted by both the Python API and the CLI. +FORMAT_ALIASES = {"gpt-oss": "gptoss"} + + +def detect_format(dataset: Dataset) -> str: + """ + Auto-detect dataset format by probing columns and turn structure. + + Returns a format name from FORMAT_REGISTRY, or 'unknown'. + Walks the registry in order; first match wins. + """ + conv_info = _probe_conversation(dataset) + for entry in FORMAT_REGISTRY: + if entry["match"](dataset, conv_info): + return entry["name"] + return "unknown" + + +def get_scanner(fmt: str): + """Return the scanner function for a format name, or None if unknown.""" + for entry in FORMAT_REGISTRY: + if entry["name"] == fmt: + return entry["scan"] + return None + + +def scan_dataset(dataset: Dataset, fmt: str = "auto") -> dict: + """ + One-liner: detect format (if 'auto') and scan for None/empty content. + + Returns the stats dict with an added 'format' key. + Raises ValueError if the format is unknown or unsupported. + """ + # Reject a DatasetDict / IterableDatasetDict (load_dataset without split): + # its column_names is a split map and would yield a confusing "unknown + # format". Check both (IterableDatasetDict is not a DatasetDict subclass); + # import locally so this module never hard-requires them. + _dict_types = [] + try: + from datasets import DatasetDict as _DatasetDict + _dict_types.append(_DatasetDict) + except ImportError: + pass + try: + from datasets import IterableDatasetDict as _IterableDatasetDict + _dict_types.append(_IterableDatasetDict) + except ImportError: + pass + if _dict_types and isinstance(dataset, tuple(_dict_types)): + raise ValueError( + "scan_dataset requires a single Dataset split, not a DatasetDict. " + f"Available splits: {list(dataset.keys())}. " + "Pass dataset[] or use load_dataset(..., split='train')." + ) + # Streaming IterableDataset has no len()/column_names; give a clear error + # instead of a confusing downstream TypeError. + try: + from datasets import IterableDataset as _IterableDataset + if isinstance(dataset, _IterableDataset): + raise ValueError( + "scan_dataset requires a materialized Dataset, not an IterableDataset. " + "Load without streaming=True, or materialize a slice first: " + "Dataset.from_list(list(dataset.take(N)))." + ) + except ImportError: + pass + fmt = FORMAT_ALIASES.get(fmt, fmt) + was_auto = fmt == "auto" + # Zero-row dataset: return a trivially clean stats dict. + if was_auto and len(dataset) == 0: + return { + "format": "unknown", + "total_rows": 0, + "findings": [], + "bad_row_indices": [], + } + # Always probe so detection and column selection share one scan pass. + conv_info = _probe_conversation(dataset) + if was_auto: + fmt = "unknown" + for entry in FORMAT_REGISTRY: + if entry["match"](dataset, conv_info): + fmt = entry["name"] + break + # No format matched: return clean stats (format="unknown") instead of + # raising, so callers can branch on stats["format"]. + if fmt == "unknown": + return { + "format": "unknown", + "total_rows": len(dataset), + "findings": [], + "bad_row_indices": [], + } + scanner = get_scanner(fmt) + if scanner is None: + raise ValueError(f"Unknown or unsupported format: '{fmt}'") + # Column forwarding: on auto-detect pass the probed column (the best + # choice). On an explicit format let that scanner pick its own column, so + # e.g. fmt='sharegpt' always scans 'conversations', not 'messages' (P1 fix); + # gptoss has its own messages-first rule. alpaca never takes a column. + use_probed_col = conv_info is not None and fmt != "alpaca" and was_auto + if use_probed_col: + stats = scanner(dataset, col = conv_info["column"]) + else: + stats = scanner(dataset) + stats["format"] = fmt + return stats + + +# --------------------------------------------------------------------------- +# Report printing +# --------------------------------------------------------------------------- + + +def _print_summary_header(stats: dict, fmt: str) -> bool: + """Print the top-level stats block (shared by all report modes). Returns True if findings exist.""" + total = stats["total_rows"] + findings = stats.get("findings", []) + + print(f"\n{'=' * 64}") + print(f" None / Empty Detection Report") + print(f"{'=' * 64}") + print(f" Format: {fmt}") + print(f" Total rows: {total}") + + if not findings: + if fmt == "unknown": + print(f" Result: NOT SCANNED -- format could not be detected") + else: + print(f" Result: CLEAN -- no None or empty values found") + print(f"{'=' * 64}") + return False + + if fmt == "alpaca": + bad_rows = len(stats.get("bad_row_indices", [])) + print(f" Rows with Nones: {bad_rows} / {total}") + print(f" None instruction: {stats.get('none_instruction', 0)}") + print(f" None output: {stats.get('none_output', 0)}") + else: + col = stats.get("column", "?") + print(f" Column: {col}") + print(f" Rows with bad turns: {stats['rows_with_none_turns']} / {total}") + print(f" Total bad turns: {len(findings)}") + print(f" By type: {stats.get('none_by_type', {})}") + print(f" By role: {stats.get('none_by_role', {})}") + rows_all = stats.get("rows_all_none", 0) + if rows_all: + print(f" Rows ALL bad: {rows_all} (every turn is None/empty)") + + # Rows with no Nones - compute the count directly rather than allocating a + # full set of row indices, which OOMs on large (10M+ row) datasets. + bad_indices = set(stats.get("bad_row_indices", [])) + clean_count = total - len(bad_indices) + if 0 < clean_count <= 20: + clean_indices = [i for i in range(total) if i not in bad_indices] + print(f" Rows with no Nones: {clean_count} / {total} {clean_indices}") + else: + print(f" Rows with no Nones: {clean_count} / {total}") + + print(f"{'=' * 64}") + return True + + +def print_report( + stats: dict, + fmt: str, + summary_only: bool = False, +): + """Print a human-readable summary, optionally with full findings list.""" + has_findings = _print_summary_header(stats, fmt) + if not has_findings or summary_only: + return + + findings = stats.get("findings", []) + print(f"\n {'-' * 60}") + print(f" Findings ({len(findings)} total):") + print(f" {'-' * 60}") + + for f in findings: + if fmt == "alpaca": + print( + f" row {f['row_index']:>5d} " + f"field={f['field']:<12s} " + f"type={f['value_type']:<16s} " + f"raw={f['raw_value']}" + ) + else: + print( + f" row {f['row_index']:>5d} " + f"turn {f['turn_index']} " + f"role={str(f['role']):<12s} " + f"type={f['value_type']:<16s} " + f"raw={f['raw_value']}" + ) + + print(f"{'=' * 64}") + + +def show_row( + dataset: Dataset, + row_indices: list[int], + fmt: str, + col: str = None, +): + """Print the full contents of specific rows for inspection. + + Used by test_codex_fixes.py to verify row rendering behaviour. + Not part of the production API. + """ + if col is None: + for candidate in ("messages", "conversations", "texts"): + if candidate in dataset.column_names: + col = candidate + break + + for ri in row_indices: + if ri < 0 or ri >= len(dataset): + print(f"\n [ERROR] Row {ri} out of range (0-{len(dataset)-1})") + continue + + row = dataset[ri] + print(f"\n{'=' * 64}") + print(f" Row {ri}") + print(f"{'=' * 64}") + + # Print non-conversation columns. For alpaca, skip fields the alpaca + # block below prints with status markers (avoid double render). + _ALPACA_FIELDS = {"instruction", "input", "output"} + for key in dataset.column_names: + if key == col: + continue + if fmt == "alpaca" and key in _ALPACA_FIELDS: + continue + val = row[key] + if isinstance(val, str) and len(val) > 120: + val = val[:120] + "..." + print(f" {key}: {val}") + + if fmt == "alpaca": + for field in ("instruction", "input", "output"): + val = row.get(field) + status = " [NONE]" if is_none_or_empty(val) else "" + if val and len(str(val)) > 200: + val = str(val)[:200] + "..." + print(f" {field}: {val}{status}") + elif col: + conversation = row[col] + if isinstance(conversation, list): + + def _is_bad_turn(t): + if not isinstance(t, dict): + return True + # Mirror scanner logic: from+value wins, then role, then from alone. + if "from" in t and "value" in t: + c = t.get("value") + elif "role" in t: + c = t.get("content") if "content" in t else t.get("value") + elif "from" in t: + c = t.get("value") + else: + c = t.get("content") if "content" in t else t.get("value") + # Mirror scanner: tool_calls exemption is assistant-only; + # other roles with empty content + tool_calls are still bad. + r = t.get("role") if t.get("role") is not None else t.get("from") + if is_none_or_empty(c) and not (str(r) == "assistant" and t.get("tool_calls")): + return True + return False + + none_count = sum(1 for t in conversation if _is_bad_turn(t)) + print(f" {col}: {len(conversation)} turns ({none_count} None)") + print(f" {'-' * 60}") + for i, turn in enumerate(conversation): + # Non-dict turn - can't extract role/content normally. + if not isinstance(turn, dict): + label = "None" if turn is None else "invalid_type" + print(f" [{i:>3d}] {'unknown':<12s} [{label}] << NONE") + continue + r = turn.get("role") + if r is None: + r = turn.get("from") + role = "?" if r is None else str(r) + # Mirror scanner logic: from+value wins, then role, then from alone. + if "from" in turn and "value" in turn: + content = turn.get("value") + elif "role" in turn: + content = turn.get("content") if "content" in turn else turn.get("value") + elif "from" in turn: + content = turn.get("value") + else: + content = turn.get("content") if "content" in turn else turn.get("value") + if is_none_or_empty(content) and not ( + role == "assistant" and turn.get("tool_calls") + ): + status = " << NONE" + else: + status = "" + if content is None: + preview = "None" + else: + preview_str = str(content) # cast: content may not be a string + if len(preview_str) > 150: + preview = preview_str[:150].replace("\n", "\\n") + "..." + else: + preview = preview_str.replace("\n", "\\n") + print(f" [{i:>3d}] {role:<12s} {preview}{status}") + + print(f"{'=' * 64}") + + +# --------------------------------------------------------------------------- +# CLI entry point +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + import argparse + import os + import sys + + parser = argparse.ArgumentParser( + prog = "dataset_none_detect", + description = "Scan a HuggingFace dataset for None/empty content turns.", + formatter_class = argparse.RawDescriptionHelpFormatter, + epilog = """ +examples: + python dataset_none_detect.py org/my-dataset + python dataset_none_detect.py org/my-dataset --split train + python dataset_none_detect.py org/my-dataset --format sharegpt + python dataset_none_detect.py org/my-dataset --summary-only + python dataset_none_detect.py org/my-dataset --token hf_... + """, + ) + parser.add_argument("dataset", help = "HuggingFace dataset repo id (e.g. org/my-dataset)") + parser.add_argument("--split", default = "train", help = "Dataset split to load (default: train)") + parser.add_argument( + "--format", + default = "auto", + choices = ["auto"] + FORMAT_NAMES + list(FORMAT_ALIASES), + help = "Force a specific format instead of auto-detecting (default: auto). " + "Documented aliases (e.g. 'gpt-oss' for 'gptoss') are also accepted.", + ) + parser.add_argument( + "--summary-only", + action = "store_true", + help = "Print summary header only - skip the per-turn findings list", + ) + parser.add_argument( + "--token", + default = os.environ.get("HF_TOKEN"), + help = ( + "HuggingFace API token for private datasets (default: $HF_TOKEN). " + "Prefer setting $HF_TOKEN; passing --token on the command line " + "exposes it in process listings." + ), + ) + args = parser.parse_args() + + try: + from datasets import load_dataset + except ImportError: + print( + "Error: 'datasets' package not found. Install with: pip install datasets", + file = sys.stderr, + ) + sys.exit(1) + + print(f"Loading {args.dataset!r} (split={args.split!r})...") + try: + ds = load_dataset(args.dataset, split = args.split, token = args.token) + except Exception as exc: + # Some `datasets` / `requests` versions include the Authorization + # header in exception messages. Redact the token before printing. + msg = str(exc) + if args.token: + msg = msg.replace(args.token, "hf_***REDACTED***") + print(f"Error loading dataset: {msg}", file = sys.stderr) + sys.exit(1) + + print(f"Loaded {len(ds)} rows, columns: {ds.column_names}") + + try: + stats = scan_dataset(ds, fmt = args.format) + except ValueError as exc: + print(f"Error: {exc}", file = sys.stderr) + sys.exit(1) + + print_report(stats, stats["format"], summary_only = args.summary_only) diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 26378d64ee..faa3deac70 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -4,12 +4,12 @@ """ Dataset utilities for format detection, conversion, and template application. -This module provides the main entry points for dataset processing: -- check_dataset_format: Lightweight check if manual mapping is needed (for frontend) -- format_dataset: Detects and normalizes dataset formats -- format_and_template_dataset: End-to-end processing with chat template application +Main entry points for dataset processing: +- check_dataset_format: lightweight check if manual mapping is needed (frontend) +- format_dataset: detects and normalizes dataset formats +- format_and_template_dataset: end-to-end processing with chat template -All internal utilities have been moved to separate modules: +Internal utilities live in separate modules: - format_detection: detect_dataset_format, detect_multimodal_dataset, etc. - format_conversion: standardize_chat_format, convert_chatml_to_alpaca, etc. - chat_templates: apply_chat_template_to_dataset, get_tokenizer_chat_template, etc. @@ -20,7 +20,6 @@ All internal utilities have been moved to separate modules: import json -# Import from modular files from .format_detection import ( detect_dataset_format, detect_multimodal_dataset, @@ -54,8 +53,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: """ Lightweight format check without processing - for frontend validation. - Use this to quickly determine if user needs to manually map columns - before calling the full format_and_template_dataset(). + Quickly determines if the user must manually map columns before the full + format_and_template_dataset(). Args: dataset: HuggingFace dataset @@ -121,7 +120,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: } if is_audio: - # Audio dataset — require manual mapping only when columns can't be auto-detected + # Audio dataset — require manual mapping only when columns aren't auto-detected detected_audio = multimodal_info.get("detected_audio_column") detected_text = multimodal_info.get("detected_text_column") needs_mapping = not detected_audio or not detected_text @@ -181,6 +180,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "suggested_mapping": None, "detected_image_column": None, "detected_text_column": None, + "chat_column": detected.get("chat_column"), "is_image": multimodal_info["is_image"], "multimodal_columns": multimodal_info.get("multimodal_columns"), **audio_fields, @@ -200,16 +200,31 @@ _TO_CHATML = { } _CHATML_ROLE_ORDER = ("system", "user", "assistant") _CHATML_TO_ALPACA = {"user": "instruction", "system": "input", "assistant": "output"} +_KNOWN_CHAT_COLUMNS = {"messages", "conversations", "texts"} -def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000): +def _chatml_final_format(chat_column: str | None) -> str: + return "chatml_messages" if chat_column == "messages" else "chatml_conversations" + + +def _chatml_detected_format_label(chat_column: str | None) -> str: + if chat_column in _KNOWN_CHAT_COLUMNS: + return f"chatml_{chat_column}" + return "chatml_conversations" + + +def _apply_user_mapping( + dataset, + mapping: dict, + batch_size: int = 1000, +): """ Apply user-provided column mapping to convert dataset to conversations format. Accepts chatml (user/assistant/system), sharegpt (human/gpt/system), and - alpaca (instruction/input/output) role names — all normalised to chatml output. + alpaca (instruction/input/output) role names — all normalised to chatml. - If the mapping contains ``__``-prefixed metadata keys (from the conversion + If the mapping has ``__``-prefixed metadata keys (from the conversion advisor), routes to template-based conversion instead of simple role mapping. Returns: @@ -258,7 +273,7 @@ def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000): def _extract_column_value(val, col: str, label_mapping: dict) -> str: """Extract a string value from a column, handling complex types and label mapping.""" - # Handle complex types (dicts, lists) — extract useful text instead of raw repr + # Complex types (dicts, lists): extract useful text instead of raw repr if isinstance(val, dict): # Common pattern: {"text": [...]} in QA datasets if "text" in val: @@ -279,15 +294,17 @@ def _extract_column_value(val, col: str, label_mapping: dict) -> str: def _apply_template_mapping( - dataset, column_roles: dict, meta: dict, batch_size: int = 1000 + dataset, + column_roles: dict, + meta: dict, + batch_size: int = 1000, ): """ Apply advisor-driven mapping for non-conversational datasets. - Groups columns by their assigned role (user/assistant), concatenates - values within each role into a single message, and injects an optional - system prompt. Label mapping is applied to convert integer labels - to human-readable strings. + Groups columns by assigned role (user/assistant), concatenates values + within each role into one message, and injects an optional system prompt. + Label mapping converts integer labels to human-readable strings. Returns: Dataset with single 'conversations' column @@ -320,23 +337,19 @@ def _apply_template_mapping( if system_prompt: convo.append({"role": "system", "content": system_prompt}) - # User message: concatenate all user-role column values + # User message: concatenate user-role column values user_parts = [] for col in role_groups["user"]: if col in examples: - user_parts.append( - _extract_column_value(examples[col][i], col, label_mapping) - ) + user_parts.append(_extract_column_value(examples[col][i], col, label_mapping)) if user_parts: convo.append({"role": "user", "content": "\n".join(user_parts)}) - # Assistant message: concatenate all assistant-role column values + # Assistant message: concatenate assistant-role column values asst_parts = [] for col in role_groups["assistant"]: if col in examples: - asst_parts.append( - _extract_column_value(examples[col][i], col, label_mapping) - ) + asst_parts.append(_extract_column_value(examples[col][i], col, label_mapping)) if asst_parts: convo.append({"role": "assistant", "content": "\n".join(asst_parts)}) @@ -351,7 +364,11 @@ def _apply_template_mapping( ) -def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000): +def _apply_user_mapping_alpaca( + dataset, + mapping: dict, + batch_size: int = 1000, +): """ Apply user-provided column mapping to convert dataset to Alpaca format. @@ -382,11 +399,7 @@ def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000): ("output", outputs), ): col = col_for[field] - val = ( - str(examples[col][i]) - if col and col in examples and examples[col][i] - else "" - ) + val = str(examples[col][i]) if col and col in examples and examples[col][i] else "" dest.append(val) return {"instruction": instructions, "input": inputs, "output": outputs} @@ -430,7 +443,7 @@ def format_dataset( "final_format": final format after processing, "chat_column": column name with chat data, "is_standardized": whether role names are standardized, - "requires_manual_mapping": True if format detection failed and user must map columns, + "requires_manual_mapping": True if detection failed and user must map columns, "warnings": list of warning messages } """ @@ -452,7 +465,7 @@ def format_dataset( "warnings": [notice.message for notice in raw_result.notices], } - # If user provided explicit mapping, skip detection and apply in the requested format + # If user provided explicit mapping, skip detection and apply it if custom_format_mapping: try: if format_type == "alpaca": @@ -462,11 +475,9 @@ def format_dataset( final_format = "alpaca" chat_column = None else: - # auto / chatml / sharegpt / conversational — all produce chatml conversations - # (sharegpt is always standardized to role/content internally) - mapped_dataset = _apply_user_mapping( - dataset, custom_format_mapping, batch_size - ) + # auto / chatml / sharegpt / conversational all produce chatml + # conversations (sharegpt standardized to role/content internally) + mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size) final_format = "chatml_conversations" chat_column = "conversations" @@ -523,7 +534,7 @@ def format_dataset( } # ShareGPT - needs standardization - elif detected["format"] == "sharegpt": + elif detected["format"] == "sharegpt" and detected.get("chat_column"): try: standardized = standardize_chat_format( dataset, @@ -533,11 +544,12 @@ def format_dataset( aliases_for_assistant, batch_size, num_proc, + chat_column = detected["chat_column"], ) return { "dataset": standardized, "detected_format": "sharegpt", - "final_format": f"chatml_{detected['chat_column']}", + "final_format": _chatml_final_format(detected["chat_column"]), "chat_column": detected["chat_column"], "is_standardized": True, "requires_manual_mapping": False, @@ -559,15 +571,11 @@ def format_dataset( "warnings": warnings, } - elif detected["format"] == "chatml" and detected["chat_column"] in [ - "conversations", - "messages", - "texts", - ]: + elif detected["format"] == "chatml" and detected.get("chat_column"): return { "dataset": dataset, - "detected_format": f"chatml_{detected['chat_column']}", - "final_format": f"chatml_{detected['chat_column']}", + "detected_format": _chatml_detected_format_label(detected["chat_column"]), + "final_format": _chatml_final_format(detected["chat_column"]), "chat_column": detected["chat_column"], "is_standardized": True, "requires_manual_mapping": False, @@ -576,13 +584,11 @@ def format_dataset( "warnings": warnings, } - # Unknown - try standardization, if fails pass as is + # Unknown - try standardization, pass as-is on failure else: - warnings.append( - f"Unknown format detected. Keys found: {detected['sample_keys']}" - ) + warnings.append(f"Unknown format detected. Keys found: {detected['sample_keys']}") - # NEW: Try heuristic detection + # Try heuristic detection if auto_detect_custom: custom_mapping = detect_custom_format_heuristic(dataset) if custom_mapping: @@ -606,9 +612,7 @@ def format_dataset( if role == target_role and col_name in examples: content = examples[col_name][i] if content and str(content).strip(): - convo.append( - {"role": role, "content": str(content)} - ) + convo.append({"role": role, "content": str(content)}) conversations.append(convo) return {"conversations": conversations, **preserved_columns} @@ -642,12 +646,13 @@ def format_dataset( aliases_for_assistant, batch_size, num_proc, + chat_column = detected["chat_column"], ) warnings.append("Successfully standardized unknown format") return { "dataset": standardized, "detected_format": "unknown", - "final_format": f"chatml_{detected['chat_column']}", + "final_format": _chatml_final_format(detected["chat_column"]), "chat_column": detected["chat_column"], "is_standardized": True, "requires_manual_mapping": False, @@ -656,9 +661,7 @@ def format_dataset( "warnings": warnings, } except Exception as e: - warnings.append( - f"Could not standardize: {e}. Passing dataset as-is." - ) + warnings.append(f"Could not standardize: {e}. Passing dataset as-is.") # Return as-is with warnings return { @@ -688,32 +691,52 @@ def format_dataset( "warnings": [], } - elif detected["format"] in ["sharegpt", "chatml"]: - # First standardize if ShareGPT - if detected["format"] == "sharegpt": - dataset = standardize_chat_format( + elif detected["format"] in ["sharegpt", "chatml"] and detected.get("chat_column"): + try: + # First standardize if ShareGPT + if detected["format"] == "sharegpt": + dataset = standardize_chat_format( + dataset, + tokenizer, + aliases_for_system, + aliases_for_user, + aliases_for_assistant, + batch_size, + num_proc, + chat_column = detected["chat_column"], + ) + + # Then convert to Alpaca + converted = convert_chatml_to_alpaca( dataset, - tokenizer, - aliases_for_system, - aliases_for_user, - aliases_for_assistant, batch_size, num_proc, + chat_column = detected["chat_column"], ) - - # Then convert to Alpaca - converted = convert_chatml_to_alpaca(dataset, batch_size, num_proc) - return { - "dataset": converted, - "detected_format": detected["format"], - "final_format": "alpaca", - "chat_column": None, - "is_standardized": True, - "requires_manual_mapping": False, - "is_image": multimodal_info["is_image"], - "multimodal_info": multimodal_info, - "warnings": [], - } + return { + "dataset": converted, + "detected_format": detected["format"], + "final_format": "alpaca", + "chat_column": None, + "is_standardized": True, + "requires_manual_mapping": False, + "is_image": multimodal_info["is_image"], + "multimodal_info": multimodal_info, + "warnings": [], + } + except Exception as e: + warnings.append(f"Failed to convert chat dataset to Alpaca: {e}") + return { + "dataset": dataset, + "detected_format": detected["format"], + "final_format": "unknown", + "chat_column": detected["chat_column"], + "is_standardized": False, + "requires_manual_mapping": True, + "is_image": multimodal_info["is_image"], + "multimodal_info": multimodal_info, + "warnings": warnings, + } else: warnings.append(f"Cannot convert unknown format to Alpaca") @@ -745,33 +768,48 @@ def format_dataset( "warnings": [], } - elif detected["format"] == "sharegpt": - standardized = standardize_chat_format( - dataset, - tokenizer, - aliases_for_system, - aliases_for_user, - aliases_for_assistant, - batch_size, - num_proc, - ) - return { - "dataset": standardized, - "detected_format": "sharegpt", - "final_format": f"chatml_{detected['chat_column']}", - "chat_column": detected["chat_column"], - "is_standardized": True, - "requires_manual_mapping": False, - "is_image": multimodal_info["is_image"], - "multimodal_info": multimodal_info, - "warnings": [], - } + elif detected["format"] == "sharegpt" and detected.get("chat_column"): + try: + standardized = standardize_chat_format( + dataset, + tokenizer, + aliases_for_system, + aliases_for_user, + aliases_for_assistant, + batch_size, + num_proc, + chat_column = detected["chat_column"], + ) + return { + "dataset": standardized, + "detected_format": "sharegpt", + "final_format": _chatml_final_format(detected["chat_column"]), + "chat_column": detected["chat_column"], + "is_standardized": True, + "requires_manual_mapping": False, + "is_image": multimodal_info["is_image"], + "multimodal_info": multimodal_info, + "warnings": [], + } + except Exception as e: + warnings.append(f"Failed to standardize ShareGPT format: {e}") + return { + "dataset": dataset, + "detected_format": "sharegpt", + "final_format": "sharegpt", + "chat_column": detected["chat_column"], + "is_standardized": False, + "requires_manual_mapping": True, + "is_image": multimodal_info["is_image"], + "multimodal_info": multimodal_info, + "warnings": warnings, + } - elif detected["format"] == "chatml": + elif detected["format"] == "chatml" and detected.get("chat_column"): return { "dataset": dataset, - "detected_format": f"chatml_{detected['chat_column']}", - "final_format": f"chatml_{detected['chat_column']}", + "detected_format": _chatml_detected_format_label(detected["chat_column"]), + "final_format": _chatml_final_format(detected["chat_column"]), "chat_column": detected["chat_column"], "is_standardized": True, "requires_manual_mapping": False, @@ -792,11 +830,12 @@ def format_dataset( aliases_for_assistant, batch_size, num_proc, + chat_column = detected["chat_column"], ) return { "dataset": standardized, "detected_format": "unknown", - "final_format": f"chatml_{detected['chat_column']}", + "final_format": _chatml_final_format(detected["chat_column"]), "chat_column": detected["chat_column"], "is_standardized": True, "requires_manual_mapping": False, @@ -858,8 +897,8 @@ def format_and_template_dataset( progress_callback = None, ): """ - Convenience function that combines format_dataset and apply_chat_template_to_dataset. - Perfect for UI workflows - one function does everything! + Combines format_dataset and apply_chat_template_to_dataset. Convenient for + UI workflows: one function does everything. Returns: dict: { @@ -867,7 +906,7 @@ def format_and_template_dataset( "detected_format": Original format, "final_format": Format after processing, "success": Whether template application succeeded, - "requires_manual_mapping": True if format detection failed and user must map columns, + "requires_manual_mapping": True if detection failed and user must map columns, "warnings": List of warnings, "errors": List of errors, "summary": Human-readable summary @@ -881,7 +920,7 @@ def format_and_template_dataset( multimodal_info = detect_multimodal_dataset(dataset) - # NEW: If user provided explicit mapping for VLM, use it directly + # If user provided explicit mapping for VLM, use it directly if custom_format_mapping: # Expect mapping like: {"image_col": "image", "caption_col": "text"} user_vlm_image_column = None @@ -921,17 +960,14 @@ def format_and_template_dataset( "errors": [], } except Exception as e: - # User mapping failed — fall back to auto-detection instead - # of giving up (handles stale cached mappings gracefully) + # User mapping failed; fall back to auto-detection (handles stale cached mappings). warnings.append( f"User VLM mapping (image='{user_vlm_image_column}', " f"text='{user_vlm_text_column}') failed: {e} — " f"falling back to auto-detection" ) - logger.info( - f"⚠️ User VLM mapping failed, falling back to auto-detection..." - ) - custom_format_mapping = None # clear so auto-detection runs below + logger.info(f"⚠️ User VLM mapping failed, falling back to auto-detection...") + custom_format_mapping = None # so auto-detection runs below else: errors.append( f"Invalid VLM mapping: need 'image' and 'text' roles. Got: {custom_format_mapping}" @@ -974,7 +1010,7 @@ def format_and_template_dataset( "errors": errors, } - # Handle ShareGPT/ChatML + image column (e.g. ShareGPT4V, LLaVA-style) + # ShareGPT/ChatML + image column (e.g. ShareGPT4V, LLaVA-style) elif vlm_structure["format"] == "sharegpt_with_images": try: dataset = convert_sharegpt_with_images_to_vlm_format( @@ -984,9 +1020,7 @@ def format_and_template_dataset( dataset_name = dataset_name, progress_callback = progress_callback, ) - warnings.append( - "Converted from ShareGPT+image format to standard VLM format" - ) + warnings.append("Converted from ShareGPT+image format to standard VLM format") except Exception as e: errors.append(f"Failed to convert ShareGPT+image format: {e}") import traceback @@ -1020,7 +1054,6 @@ def format_and_template_dataset( friendly = None try: from .llm_assist import llm_generate_dataset_warning - friendly = llm_generate_dataset_warning( issues, dataset_name = dataset_name, @@ -1055,13 +1088,9 @@ def format_and_template_dataset( ) if vlm_instruction: - warnings.append( - f"Using user-provided instruction: '{vlm_instruction}'" - ) + warnings.append(f"Using user-provided instruction: '{vlm_instruction}'") else: - warnings.append( - "Auto-generated instruction based on dataset analysis" - ) + warnings.append("Auto-generated instruction based on dataset analysis") except Exception as e: errors.append(f"Failed to convert to VLM format: {e}") @@ -1101,7 +1130,7 @@ def format_and_template_dataset( "errors": errors, } - # LLM FLOW (Existing code) + # LLM FLOW else: # Step 1: Format the dataset n_rows = len(dataset) if hasattr(dataset, "__len__") else None @@ -1141,7 +1170,7 @@ def format_and_template_dataset( progress_callback( status_message = f"Applying chat template to {detected} ({n_rows:,} rows)..." ) - # Gemma emits a leading that must be stripped for text-only chatml/sharegpt. + # Gemma emits a leading , stripped for text-only chatml/sharegpt. is_alpaca = format_type == "alpaca" or ( format_type == "auto" and dataset_info["detected_format"] == "alpaca" ) @@ -1166,13 +1195,10 @@ def format_and_template_dataset( summary = get_dataset_info_summary(dataset_info) # Combine results - all_warnings = dataset_info.get("warnings", []) + template_result.get( - "warnings", [] - ) + all_warnings = dataset_info.get("warnings", []) + template_result.get("warnings", []) all_errors = template_result.get("errors", []) - # If format_dataset returned "unknown" but apply_chat_template rescued - # it via heuristic detection, update final_format to reflect reality. + # If apply_chat_template rescued an "unknown" format, update final_format. final_format = dataset_info["final_format"] requires_manual = dataset_info.get("requires_manual_mapping", False) if final_format == "unknown" and template_result["success"]: @@ -1186,7 +1212,7 @@ def format_and_template_dataset( "detected_format": dataset_info["detected_format"], "final_format": final_format, "chat_column": dataset_info.get("chat_column"), - "is_vlm": False, # This is LLM flow + "is_vlm": False, # LLM flow "success": template_result["success"], "requires_manual_mapping": requires_manual, "warnings": all_warnings, diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index 289b30e55e..cb24bd96ba 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -1,12 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Format conversion utilities for dataset processing. - -This module contains functions for converting between dataset formats -(Alpaca, ShareGPT, ChatML) and standardizing chat formats. -""" +"""Dataset format conversion between Alpaca, ShareGPT, and ChatML.""" import os @@ -34,16 +29,17 @@ def standardize_chat_format( ], batch_size = 1000, num_proc = None, + chat_column: str | None = None, ): """ - Our own standardization function that handles BOTH messages and conversations. - Converts non-standard role names and keys to standard format. + Standardize BOTH messages and conversations: map non-standard role + names and keys to the standard format. """ import collections import itertools from datasets import IterableDataset - # Check if vision tokenizer is used + # Detect a vision tokenizer is_vlm = False if tokenizer is not None: if hasattr(tokenizer, "image_processor") or hasattr(tokenizer, "tokenizer"): @@ -51,9 +47,10 @@ def standardize_chat_format( column_names = set(next(iter(dataset)).keys()) - # Check for both 'conversations' and 'messages' - chat_column = None - if "conversations" in column_names: + if chat_column: + if chat_column not in column_names: + return dataset + elif "conversations" in column_names: chat_column = "conversations" elif "messages" in column_names: chat_column = "messages" @@ -62,30 +59,48 @@ def standardize_chat_format( else: return dataset # No chat column found - # Inspect structure - examples = itertools.islice(dataset, 10) + def _iter_probe_rows(): + try: + total = min(len(dataset), 100) + for index in range(total): + yield dataset[index] + return + except Exception: + pass + for example in itertools.islice(dataset, 100): + yield example + uniques = collections.defaultdict(list) - for example in examples: - for message in example[chat_column]: + for example in _iter_probe_rows(): + chat_data = example.get(chat_column) + if not isinstance(chat_data, list) or len(chat_data) == 0: + continue + for message in chat_data: + if not isinstance(message, dict): + continue for key, value in message.items(): if type(value) is not str: - continue # Skip non-string values + continue # Skip non-strings uniques[key].append(value) - if len(uniques.keys()) != 2: - return dataset # Unexpected structure - - keys = list(uniques.keys()) - length_first = len(set(uniques[keys[0]])) - length_second = len(set(uniques[keys[1]])) - - # Determine which is role and which is content - if length_first < length_second: - role_key = keys[0] - content_key = keys[1] + if "from" in uniques and "value" in uniques: + role_key = "from" + content_key = "value" + elif "role" in uniques and "content" in uniques: + role_key = "role" + content_key = "content" + elif len(uniques.keys()) == 2: + keys = list(uniques.keys()) + length_first = len(set(uniques[keys[0]])) + length_second = len(set(uniques[keys[1]])) + if length_first < length_second: + role_key = keys[0] + content_key = keys[1] + else: + role_key = keys[1] + content_key = keys[0] else: - role_key = keys[1] - content_key = keys[0] + raise ValueError(f"Could not infer role/content keys for chat column '{chat_column}'") # Mapping for aliases aliases_mapping = {} @@ -100,20 +115,30 @@ def standardize_chat_format( convos = examples[chat_column] all_convos = [] for convo in convos: + if not isinstance(convo, list): + all_convos.append([]) + continue + new_convo = [] for message in convo: - # Get original role and content - original_role = message.get(role_key, "") - original_content = message.get(content_key, "") + if not isinstance(message, dict): + continue + + # Use the inferred keys first; fall back per-message so mixed + # ShareGPT/ChatML rows keep valid turns. + original_role = message.get(role_key) + original_content = message.get(content_key) + if original_role is None: + original_role = message.get("role") or message.get("from") or "" + if original_content is None: + original_content = message.get("content") or message.get("value") or "" - # Map to standard role name standard_role = aliases_mapping.get(original_role, original_role) - # Handle VLM format if is_vlm: original_content = [{"type": "text", "text": original_content}] - # Create dict with EXPLICIT ORDER + # Keep EXPLICIT key order new_message = {"role": standard_role, "content": original_content} new_convo.append(new_message) @@ -140,10 +165,14 @@ def standardize_chat_format( return dataset.map(_standardize_dataset, **dataset_map_kwargs) -def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None): +def convert_chatml_to_alpaca( + dataset, + batch_size = 1000, + num_proc = None, + chat_column: str | None = None, +): """ - Converts ChatML format (messages OR conversations) to Alpaca format. - Handles both standardized and ShareGPT formats. + Convert ChatML (messages OR conversations) to Alpaca format. Supports: - "messages" or "conversations" column @@ -151,23 +180,19 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None): """ try: from torch.utils.data import IterableDataset - _is_torch_iterable = isinstance(dataset, IterableDataset) except ImportError: _is_torch_iterable = False def _convert(examples): - # Auto-detect which column name is used - chatml_data = ( - examples.get("messages") - or examples.get("conversations") - or examples.get("texts") - ) + chatml_data = examples.get(chat_column) if chat_column else None + if chatml_data is None: + chatml_data = ( + examples.get("messages") or examples.get("conversations") or examples.get("texts") + ) if chatml_data is None: - raise ValueError( - "No 'messages' or 'conversations' or 'texts' column found." - ) + raise ValueError("No 'messages' or 'conversations' or 'texts' column found.") instructions = [] outputs = [] @@ -178,20 +203,20 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None): output = "" for msg in convo: - # Handle both standard and ShareGPT formats + # Standard and ShareGPT key names role = msg.get("role") or msg.get("from") content = msg.get("content") or msg.get("value") - # Get first user message as instruction + # First user message -> instruction if role in ["user", "human", "input"] and not instruction: instruction = content - # Get first assistant message as output + # First assistant message -> output elif role in ["assistant", "gpt", "output"] and not output: output = content break # Stop after first assistant response instructions.append(instruction) - inputs.append("") # Alpaca typically has empty input + inputs.append("") # Alpaca input usually empty outputs.append(output) return {"instruction": instructions, "input": inputs, "output": outputs} @@ -215,15 +240,18 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None): return dataset.map(_convert, **dataset_map_kwargs) -def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None): +def convert_alpaca_to_chatml( + dataset, + batch_size = 1000, + num_proc = None, +): """ - Converts Alpaca format to ChatML format. + Convert Alpaca format to ChatML format. - Output format: Uses 'conversations' column with standard 'role'/'content' structure. + Output: 'conversations' column with standard 'role'/'content' dicts. """ try: from torch.utils.data import IterableDataset - _is_torch_iterable = isinstance(dataset, IterableDataset) except ImportError: _is_torch_iterable = False @@ -236,13 +264,12 @@ def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None): input_text = examples.get("input", [""] * len(examples["instruction"]))[i] output = examples["output"][i] - # Combine instruction and input (if exists) for user message + # User message = instruction + input (if any) if input_text and input_text.strip(): user_content = f"{instruction}\n\n{input_text}".strip() else: user_content = instruction - # Build conversation in standard ChatML format convo = [ {"role": "user", "content": user_content}, {"role": "assistant", "content": output}, @@ -292,17 +319,14 @@ def convert_to_vlm_format( progress_callback = None, ): """ - Converts simple {image, text} format to VLM messages format. + Convert simple {image, text} format to VLM messages format. Returns a LIST, not a HuggingFace Dataset (to preserve PIL Images). - - For URL-based image datasets, runs a 200-sample parallel probe first to - estimate download speed and failure rate, then reports time estimate or - warning through progress_callback before proceeding with the full conversion. + For URL-based datasets, runs a 200-sample parallel probe first to + estimate speed/failure rate via progress_callback. Args: - progress_callback: Optional callable(status_message=str) to report - progress to the training overlay. + progress_callback: Optional callable(status_message=str) for progress. Returns: list: List of dicts with 'messages' field @@ -311,11 +335,11 @@ def convert_to_vlm_format( from .vlm_processing import generate_smart_vlm_instruction def _notify(msg): - """Send status update to the training overlay if callback is available.""" + """Send a status update to the training overlay if callback set.""" if progress_callback: progress_callback(status_message = msg) - # Generate smart instruction if not provided + # Generate a smart instruction if none provided if instruction is None: instruction_info = generate_smart_vlm_instruction( dataset, @@ -328,36 +352,30 @@ def convert_to_vlm_format( instruction_column = instruction_info.get("instruction_column") uses_dynamic = instruction_info["uses_dynamic_instruction"] - logger.info( - f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}" - ) + logger.info(f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}") logger.info(f"📝 Confidence: {instruction_info['confidence']:.2f}") if not uses_dynamic: logger.info(f"📝 Using instruction: '{instruction}'") else: - logger.info( - f"📝 Using dynamic instructions from column: '{instruction_column}'" - ) + logger.info(f"📝 Using dynamic instructions from column: '{instruction_column}'") else: instruction_column = None uses_dynamic = False def _convert_single_sample(sample): """Convert a single sample to VLM format.""" - # Get image (might be PIL Image, local path, URL, or bare filename) + # Image may be a PIL Image, local path, URL, or bare filename image_data = sample[image_column] if isinstance(image_data, str): if image_data.startswith(("http://", "https://")): import fsspec from io import BytesIO - with fsspec.open(image_data, "rb", expand = True) as f: image_data = Image.open(BytesIO(f.read())).convert("RGB") elif _image_lookup is not None and image_data in _image_lookup: # Bare filename → resolve via HF repo lookup from huggingface_hub import hf_hub_download - local_path = hf_hub_download( dataset_name, _image_lookup[image_data], @@ -367,20 +385,18 @@ def convert_to_vlm_format( else: image_data = Image.open(image_data).convert("RGB") - # Get text (if list of strings, pick a random one — e.g. multiple captions) + # Text: if a list (e.g. multiple captions), pick one at random text_data = sample[text_column] if isinstance(text_data, list) and len(text_data) > 0: import random - text_data = random.choice(text_data) - # Get instruction (static or dynamic) + # Instruction: static or dynamic if uses_dynamic and instruction_column: current_instruction = sample[instruction_column] else: current_instruction = instruction - # Build VLM messages - simple structure messages = [ { "role": "user", @@ -392,18 +408,14 @@ def convert_to_vlm_format( {"role": "assistant", "content": [{"type": "text", "text": text_data}]}, ] - # Return dict with messages return {"messages": messages} total = len(dataset) first_image = next(iter(dataset))[image_column] - has_urls = isinstance(first_image, str) and first_image.startswith( - ("http://", "https://") - ) + has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://")) - # ── Bare-filename detection: images stored as filenames (e.g. "img_001.png") - # that don't exist locally. Build a basename→repo_path lookup so we can - # resolve them via hf_hub_download during conversion. + # ── Bare-filename detection: build a basename→repo_path lookup so + # filename-only images resolve via hf_hub_download during conversion. _image_lookup = None _IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff") if ( @@ -438,7 +450,7 @@ def convert_to_vlm_format( logger.info(f"⚠️ Failed to build HF repo image lookup: {e}") _image_lookup = None - # ── URL probe: 200 samples with parallel workers to estimate speed + failure rate ── + # ── URL probe: 200 parallel samples to estimate speed + failure rate ── PROBE_SIZE = 200 MAX_FAIL_RATE = 0.3 @@ -449,9 +461,7 @@ def convert_to_vlm_format( num_workers = safe_thread_num_proc() _notify(f"Probing {PROBE_SIZE} image URLs with {num_workers} workers...") - logger.info( - f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers..." - ) + logger.info(f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers...") probe_samples = [dataset[i] for i in range(PROBE_SIZE)] probe_ok = 0 @@ -459,9 +469,7 @@ def convert_to_vlm_format( probe_start = time.time() with ThreadPoolExecutor(max_workers = num_workers) as executor: - futures = { - executor.submit(_convert_single_sample, s): s for s in probe_samples - } + futures = {executor.submit(_convert_single_sample, s): s for s in probe_samples} for future in as_completed(futures): try: future.result() @@ -479,11 +487,10 @@ def convert_to_vlm_format( f"{fail_rate:.0%} of the first {PROBE_SIZE} image URLs failed to download ({probe_fail}/{probe_total})", "Images are external URLs, not embedded in the dataset", ] - # Try LLM-friendly warning + # LLM-friendly warning friendly = None try: from .llm_assist import llm_generate_dataset_warning - friendly = llm_generate_dataset_warning( issues, dataset_name = dataset_name, @@ -502,7 +509,7 @@ def convert_to_vlm_format( _notify(msg) raise ValueError(msg) - # Estimate total time for remaining samples + # Estimate time for remaining samples remaining = total - PROBE_SIZE estimated_seconds = remaining / throughput if throughput > 0 else 0 eta_str = _format_eta(estimated_seconds) @@ -554,13 +561,11 @@ def convert_to_vlm_format( except Exception as e: failed_count += 1 if failed_count == 1: - logger.info( - f"First VLM conversion failure: {type(e).__name__}: {e}" - ) + logger.info(f"First VLM conversion failure: {type(e).__name__}: {e}") converted_list.extend(r for r in batch_results if r is not None) - # Progress update every batch + # Per-batch progress update elapsed = time.time() - start_time done = batch_end rate = done / elapsed if elapsed > 0 else 0 @@ -572,7 +577,7 @@ def convert_to_vlm_format( ) _notify(progress_msg) else: - # Sequential conversion for local/embedded images (fast, no I/O bottleneck) + # Sequential conversion for local/embedded images (no I/O bottleneck) pbar = tqdm(dataset, total = total, desc = "Converting VLM samples", unit = "sample") for sample in pbar: try: @@ -581,9 +586,7 @@ def convert_to_vlm_format( failed_count += 1 if failed_count == 1: # Log the first failure to aid debugging - logger.info( - f"First VLM conversion failure: {type(e).__name__}: {e}" - ) + logger.info(f"First VLM conversion failure: {type(e).__name__}: {e}") pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False) pbar.close() @@ -592,7 +595,7 @@ def convert_to_vlm_format( logger.info( f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images" ) - # For datasets that skipped the probe (small URL datasets), check fail rate now + # Small URL datasets skip the probe; check fail rate here if has_urls and fail_rate >= MAX_FAIL_RATE: issues = [ f"{fail_rate:.0%} of images failed to download ({failed_count}/{total})", @@ -601,7 +604,6 @@ def convert_to_vlm_format( friendly = None try: from .llm_assist import llm_generate_dataset_warning - friendly = llm_generate_dataset_warning( issues, dataset_name = dataset_name, @@ -627,7 +629,6 @@ def convert_to_vlm_format( friendly = None try: from .llm_assist import llm_generate_dataset_warning - friendly = llm_generate_dataset_warning( issues, dataset_name = dataset_name, @@ -647,7 +648,7 @@ def convert_to_vlm_format( logger.info(f"✅ Converted {len(converted_list)}/{total} samples") _notify(f"Converted {len(converted_list):,}/{total:,} images successfully") - # Return list, NOT Dataset + # Return list, NOT a Dataset return converted_list @@ -659,8 +660,8 @@ def convert_sharegpt_with_images_to_vlm_format( progress_callback = None, ): """ - Converts ShareGPT/ChatML datasets that have a separate image column and - ```` placeholders inside the conversation text. + Convert ShareGPT/ChatML datasets with a separate image column and + ```` placeholders in the conversation text. Example input:: @@ -690,7 +691,7 @@ def convert_sharegpt_with_images_to_vlm_format( if progress_callback: progress_callback(status_message = msg) - # ── Resolve image loading strategy (same 3-tier as convert_to_vlm_format) ── + # ── Resolve image loading (same 3-tier as convert_to_vlm_format) ── total = len(dataset) first_image = next(iter(dataset))[image_column] @@ -714,7 +715,7 @@ def convert_sharegpt_with_images_to_vlm_format( for f in repo_files if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS) } - # Also add the full relative paths as keys (for paths like "sam/images/sa_545504.jpg") + # Also key by full relative path (e.g. "sam/images/sa_545504.jpg") for f in repo_files: if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS): _image_lookup[f] = f @@ -732,19 +733,17 @@ def convert_sharegpt_with_images_to_vlm_format( _image_lookup = None def _resolve_image(image_data): - """Resolve image data to a PIL Image object.""" + """Resolve image data to a PIL Image.""" if hasattr(image_data, "size") and hasattr(image_data, "mode"): return image_data # Already PIL if isinstance(image_data, str): if image_data.startswith(("http://", "https://")): import fsspec from io import BytesIO - with fsspec.open(image_data, "rb", expand = True) as f: return Image.open(BytesIO(f.read())).convert("RGB") elif _image_lookup is not None and image_data in _image_lookup: from huggingface_hub import hf_hub_download - local_path = hf_hub_download( dataset_name, _image_lookup[image_data], @@ -753,19 +752,16 @@ def convert_sharegpt_with_images_to_vlm_format( return Image.open(local_path).convert("RGB") else: return Image.open(image_data).convert("RGB") - if isinstance(image_data, dict) and ( - "bytes" in image_data or "path" in image_data - ): + if isinstance(image_data, dict) and ("bytes" in image_data or "path" in image_data): if image_data.get("bytes"): from io import BytesIO - return Image.open(BytesIO(image_data["bytes"])).convert("RGB") if image_data.get("path"): return Image.open(image_data["path"]).convert("RGB") raise ValueError(f"Cannot resolve image: {type(image_data)}") def _convert_single_sample(sample): - """Convert a single ShareGPT+image sample to standard VLM format.""" + """Convert one ShareGPT+image sample to standard VLM format.""" pil_image = _resolve_image(sample[image_column]) conversation = sample[messages_column] @@ -775,7 +771,7 @@ def convert_sharegpt_with_images_to_vlm_format( role = _ROLE_MAP.get(role_raw.lower(), role_raw.lower()) text = msg.get("value") or msg.get("content") or "" - # Split on to interleave text and image content blocks + # Interleave text and image blocks around if "" in text: parts = text.split("") content = [] @@ -785,7 +781,7 @@ def convert_sharegpt_with_images_to_vlm_format( content.append({"type": "text", "text": part}) if i < len(parts) - 1: content.append({"type": "image", "image": pil_image}) - # If was the entire text, content might just be the image + # If text was only , content is just the image if not content: content.append({"type": "image", "image": pil_image}) else: @@ -812,9 +808,7 @@ def convert_sharegpt_with_images_to_vlm_format( pbar.close() if failed_count > 0: - logger.info( - f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples" - ) + logger.info(f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples") if len(converted_list) == 0: raise ValueError( @@ -829,7 +823,7 @@ def convert_sharegpt_with_images_to_vlm_format( def convert_llava_to_vlm_format(dataset): """ - Converts Llava format to standard VLM format. + Convert Llava format to standard VLM format. Llava format: - messages: [{'content': [{'type': 'image', 'index': 0}, {'type': 'text', 'text': '...'}]}] @@ -840,28 +834,25 @@ def convert_llava_to_vlm_format(dataset): """ from PIL import Image - logger.info( - f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format..." - ) + logger.info(f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format...") def _convert_single_sample(sample): - """Convert a single llava sample to standard VLM format.""" + """Convert one llava sample to standard VLM format.""" messages = sample["messages"] images = sample.get("images", []) - # Process each message new_messages = [] for msg in messages: new_content = [] for item in msg["content"]: if item["type"] == "image": - # Replace index with actual PIL image + # Replace index with the actual PIL image if "index" in item and item["index"] is not None: img_idx = item["index"] if img_idx < len(images): pil_image = images[img_idx] - # Ensure it's PIL + # Ensure PIL if isinstance(pil_image, str): pil_image = Image.open(pil_image).convert("RGB") @@ -872,7 +863,7 @@ def convert_llava_to_vlm_format(dataset): } ) else: - # No index, try to use first image + # No index: use the first image if len(images) > 0: pil_image = images[0] if isinstance(pil_image, str): @@ -881,14 +872,12 @@ def convert_llava_to_vlm_format(dataset): new_content.append({"type": "image", "image": pil_image}) elif item["type"] == "text": - # Keep text as-is (only type + text) new_content.append({"type": "text", "text": item.get("text", "")}) new_messages.append({"role": msg["role"], "content": new_content}) return {"messages": new_messages} - # Convert using list comprehension converted_list = [_convert_single_sample(sample) for sample in dataset] logger.info(f"✅ Converted {len(converted_list)} samples") diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py index 7b70ff3a76..f5ea5ca138 100644 --- a/studio/backend/utils/datasets/format_detection.py +++ b/studio/backend/utils/datasets/format_detection.py @@ -1,41 +1,156 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Format detection utilities for dataset processing. - -This module contains functions for detecting dataset formats (Alpaca, ShareGPT, ChatML), -detecting multimodal/VLM dataset structures, and heuristic-based column mapping. -""" +"""Dataset format detection: Alpaca/ShareGPT/ChatML, multimodal/VLM structures, heuristic column mapping.""" import re def _keyword_in_column(keyword: str, col_name: str) -> bool: """Word-boundary keyword match to avoid false positives like 'pic' in 'topic'.""" - return ( - re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) - is not None - ) + return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None + + +CONVERSATION_COLUMNS = ("messages", "conversations", "texts") +_CHATML_KEYS = frozenset({"role", "content"}) +_SHAREGPT_KEYS = frozenset({"from", "value"}) +_TRACE_SUFFIXES = ("__trace", "_trace") + + +def _sample_dataset_rows(dataset, limit: int = 100) -> list[dict]: + try: + total = min(len(dataset), limit) + return [dataset[index] for index in range(total)] + except Exception: + rows = [] + try: + for index, row in enumerate(dataset): + if index >= limit: + break + rows.append(row) + except Exception: + return [] + return rows + + +def _get_dataset_column_names(dataset, sample: dict) -> list[str]: + column_names = getattr(dataset, "column_names", None) + if isinstance(column_names, list): + return [str(column) for column in column_names] + return [str(column) for column in sample.keys()] + + +def _is_trace_conversation_name(column_name: str) -> bool: + return column_name.lower().endswith(_TRACE_SUFFIXES) + + +def _inspect_conversation_column(rows: list[dict], column_name: str) -> dict | None: + turn_keys: set[str] = set() + has_chatml = False + has_sharegpt = False + + for row in rows: + if not isinstance(row, dict) or column_name not in row: + continue + chat_data = row[column_name] + if not isinstance(chat_data, list) or len(chat_data) == 0: + continue + for turn in chat_data: + if not isinstance(turn, dict): + continue + keys = {str(key) for key in turn.keys()} + turn_keys.update(keys) + if _SHAREGPT_KEYS.issubset(keys): + has_sharegpt = True + if _CHATML_KEYS.issubset(keys): + has_chatml = True + + if has_sharegpt: + return { + "format": "sharegpt", + "chat_column": column_name, + "needs_standardization": True, + "sample_keys": sorted(turn_keys), + } + if has_chatml: + return { + "format": "chatml", + "chat_column": column_name, + "needs_standardization": False, + "sample_keys": sorted(turn_keys), + } + if turn_keys: + return { + "format": "unknown", + "chat_column": column_name, + "needs_standardization": None, + "sample_keys": sorted(turn_keys), + } + return None + + +def _detect_conversation_column(rows: list[dict], column_names: list[str]) -> dict | None: + column_name_set = set(column_names) + unknown_exact = None + for column_name in CONVERSATION_COLUMNS: + if column_name not in column_name_set: + continue + inspected = _inspect_conversation_column(rows, column_name) + if inspected and inspected["format"] in {"sharegpt", "chatml"}: + return inspected + if inspected and unknown_exact is None: + unknown_exact = inspected + + structural_candidates = [] + for column_name in column_names: + if column_name in CONVERSATION_COLUMNS: + continue + inspected = _inspect_conversation_column(rows, column_name) + if inspected and inspected["format"] in {"sharegpt", "chatml"}: + structural_candidates.append(inspected) + + trace_candidates = [ + candidate + for candidate in structural_candidates + if _is_trace_conversation_name(candidate["chat_column"]) + ] + if len(trace_candidates) == 1: + return trace_candidates[0] + if len(trace_candidates) > 1: + return unknown_exact + if len(structural_candidates) == 1: + return structural_candidates[0] + if unknown_exact is not None: + return unknown_exact + return None def detect_dataset_format(dataset): - """ - Detects dataset format by inspecting structure. + """Detect dataset format by inspecting structure. Returns: dict: { "format": "alpaca" | "sharegpt" | "chatml" | "unknown", - "chat_column": "messages" | "conversations" | None, + "chat_column": str | None, "needs_standardization": bool, "sample_keys": list of keys found in messages (for debugging) } """ - column_names = set(next(iter(dataset)).keys()) + sample_rows = _sample_dataset_rows(dataset) + if not sample_rows: + return { + "format": "unknown", + "chat_column": None, + "needs_standardization": None, + "sample_keys": [], + } - # Check for Alpaca + column_names = _get_dataset_column_names(dataset, sample_rows[0]) + column_name_set = set(column_names) + + # Alpaca alpaca_columns = {"instruction", "output"} - if alpaca_columns.issubset(column_names): + if alpaca_columns.issubset(column_name_set): return { "format": "alpaca", "chat_column": None, @@ -43,61 +158,10 @@ def detect_dataset_format(dataset): "sample_keys": [], } - # Check for chat-based formats (messages or conversations) - chat_column = None - if "messages" in column_names: - chat_column = "messages" - elif "conversations" in column_names: - chat_column = "conversations" - elif "texts" in column_names: - chat_column = "texts" + conversation = _detect_conversation_column(sample_rows, column_names) + if conversation: + return conversation - if chat_column: - # Inspect the structure to determine if ShareGPT or ChatML - try: - sample = next(iter(dataset)) - chat_data = sample[chat_column] - - if chat_data and len(chat_data) > 0: - first_msg = chat_data[0] - msg_keys = set(first_msg.keys()) - - # ShareGPT uses "from" and "value" - if "from" in msg_keys or "value" in msg_keys: - return { - "format": "sharegpt", - "chat_column": chat_column, - "needs_standardization": True, - "sample_keys": list(msg_keys), - } - - # ChatML uses "role" and "content" - elif "role" in msg_keys and "content" in msg_keys: - return { - "format": "chatml", - "chat_column": chat_column, - "needs_standardization": False, - "sample_keys": list(msg_keys), - } - - # Unknown structure but has chat column - else: - return { - "format": "unknown", - "chat_column": chat_column, - "needs_standardization": None, - "sample_keys": list(msg_keys), - } - except Exception as e: - return { - "format": "unknown", - "chat_column": chat_column, - "needs_standardization": None, - "sample_keys": [], - "error": str(e), - } - - # No recognized format return { "format": "unknown", "chat_column": None, @@ -107,8 +171,7 @@ def detect_dataset_format(dataset): def detect_custom_format_heuristic(dataset): - """ - Smart detection with priority scoring. + """Detection with priority scoring. Strategy for ambiguous keywords like 'task': 1. Detect assistant first (unambiguous) @@ -121,7 +184,6 @@ def detect_custom_format_heuristic(dataset): mapping = {} - # Keywords assistant_words = [ "output", "answer", @@ -138,7 +200,6 @@ def detect_custom_format_heuristic(dataset): "solve", ] - # Split into high/low priority user_words_high_priority = [ "input", "question", @@ -162,10 +223,10 @@ def detect_custom_format_heuristic(dataset): "persona", "role", "template", - "task", # Also in system + "task", # also a system keyword ] - # Metadata columns to ignore + # Metadata columns to ignore. metadata_exact_match = { "id", "idx", @@ -200,7 +261,7 @@ def detect_custom_format_heuristic(dataset): } def has_keyword(col_name, keywords): - """Check if any keyword appears in column name.""" + """True if any keyword appears in the column name.""" col_lower = col_name.lower() col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "") @@ -210,7 +271,7 @@ def detect_custom_format_heuristic(dataset): return False def is_metadata(col_name): - """Check if column is likely metadata.""" + """True if the column is likely metadata.""" col_lower = col_name.lower() if col_lower in metadata_exact_match: @@ -220,10 +281,7 @@ def detect_custom_format_heuristic(dataset): return True for pattern in metadata_prefix_patterns: - if ( - col_lower.startswith(pattern.split("_")[0] + "_") - and col_lower != pattern - ): + if col_lower.startswith(pattern.split("_")[0] + "_") and col_lower != pattern: if "_" in col_lower: prefix = col_lower.split("_")[0] if prefix in ["generation", "pass", "inference"]: @@ -235,7 +293,7 @@ def detect_custom_format_heuristic(dataset): return False def get_priority_score(col_name): - """Calculate priority score based on column name patterns.""" + """Priority score from column-name patterns.""" col_lower = col_name.lower() score = 0 @@ -246,7 +304,7 @@ def detect_custom_format_heuristic(dataset): return score def get_content_length(col_name): - """Get average content length for this column.""" + """Average content length for this column.""" try: if col_name in sample and sample[col_name]: content = str(sample[col_name]) @@ -256,21 +314,18 @@ def detect_custom_format_heuristic(dataset): return 0 def score_column(col_name, keywords, role_type, num_candidates): - """Score a column for how likely it is to be a particular role.""" + """Score how likely a column is to be a given role.""" if not has_keyword(col_name, keywords): return 0 score = 0 score += 10 - # Penalize ambiguous keywords when scoring for user + # Penalize ambiguous "task" so other user columns win. if role_type == "user": col_lower = col_name.lower() - # If column is ONLY "task" (or task_xxx), give it lower priority for user role - if "task" in col_lower and not any( - kw in col_lower for kw in user_words_high_priority - ): - score -= 15 # Significant penalty so other user columns win + if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority): + score -= 15 priority_bonus = get_priority_score(col_name) score += priority_bonus @@ -297,21 +352,15 @@ def detect_custom_format_heuristic(dataset): return score - # Filter out metadata columns content_columns = [col for col in all_columns if not is_metadata(col)] - # Count candidates first - assistant_potential = [ - col for col in content_columns if has_keyword(col, assistant_words) - ] + assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)] user_potential = [col for col in content_columns if has_keyword(col, user_words)] - # STEP 1: Find best ASSISTANT column + # STEP 1: best ASSISTANT column assistant_candidates = [] for col in assistant_potential: - score = score_column( - col, assistant_words, "assistant", len(assistant_potential) - ) + score = score_column(col, assistant_words, "assistant", len(assistant_potential)) if score > 0: assistant_candidates.append((col, score)) @@ -322,7 +371,7 @@ def detect_custom_format_heuristic(dataset): else: assistant_col = None - # STEP 2: Find best USER column (with penalty for ambiguous keywords) + # STEP 2: best USER column (penalizing ambiguous keywords) user_candidates = [] for col in user_potential: if col == assistant_col: @@ -338,35 +387,32 @@ def detect_custom_format_heuristic(dataset): else: user_col = None - # STEP 3: Check ALL remaining columns for SYSTEM matches (priority check) + # STEP 3: check remaining columns for SYSTEM matches remaining_columns = [col for col in content_columns if col not in mapping] system_col = None for col in remaining_columns: if has_keyword(col, system_words): - # Found a system match in remaining columns mapping[col] = "system" system_col = col break - # STEP 4: Handle any additional remaining columns + # STEP 4: handle any additional remaining columns if system_col: remaining_columns = [col for col in remaining_columns if col != system_col] if len(remaining_columns) >= 1: remaining_col = remaining_columns[0] - # If no strong keyword match, decide based on what's missing + # No strong keyword match: decide by what's missing. if not has_keyword(remaining_col, user_words + assistant_words): mapping[remaining_col] = "system" elif user_col is None: - # No user column yet, assign this as user mapping[remaining_col] = "user" else: - # Already have user + assistant, treat as system context mapping[remaining_col] = "system" - # VALIDATION: Ensure we have at least user + assistant + # Ensure at least user + assistant. has_user = any(role == "user" for role in mapping.values()) has_assistant = any(role == "assistant" for role in mapping.values()) @@ -384,28 +430,15 @@ def detect_custom_format_heuristic(dataset): def detect_multimodal_dataset(dataset): - """ - Detects if dataset contains multimodal data (images and/or audio). + """Detect multimodal data (images and/or audio) in a dataset. - Two-pass approach for each modality: - 1. Column-name heuristic (fast): checks for keywords. - 2. Value-type inspection (reliable): checks actual sample values. - - Returns: - dict: { - "is_image": bool, - "multimodal_columns": list of column names containing image data, - "modality_types": list of detected types (e.g., ["image", "audio"]), - "is_audio": bool, - "audio_columns": list of column names containing audio data, - "detected_audio_column": str or None, - "detected_text_column": str or None, - } + Two passes per modality: column-name keyword heuristic, then value-type + inspection. Returns a dict with is_image/is_audio flags, detected columns, + modality types, and detected audio/text/speaker columns. """ sample = next(iter(dataset)) column_names = list(sample.keys()) - # Keywords that indicate image data image_keywords = [ "image", "img", @@ -426,7 +459,6 @@ def detect_multimodal_dataset(dataset): "filename", ] - # Keywords that indicate audio data audio_keywords = ["audio", "speech", "wav", "waveform", "sound"] multimodal_columns = [] @@ -434,8 +466,7 @@ def detect_multimodal_dataset(dataset): modality_types = set() # ── Image detection ───────────────────────────────────── - # Pass 1: column-name heuristic (word-boundary match to avoid - # false positives like 'pic' in 'topic') + # Pass 1: column-name heuristic (word-boundary match) for col_name in column_names: for keyword in image_keywords: if _keyword_in_column(keyword, col_name): @@ -472,13 +503,13 @@ def detect_multimodal_dataset(dataset): audio_columns.append(col_name) modality_types.add("audio") - # Filter out columns that are actually audio from the image list - # (e.g. a column named "audio" with {"bytes", "path"} could match _is_image_value) + # Drop audio columns from the image list (a {"bytes","path"} audio column + # can match _is_image_value). if audio_columns: audio_set = set(audio_columns) multimodal_columns = [c for c in multimodal_columns if c not in audio_set] - # Detect text column for audio datasets + # Text column for audio datasets. detected_text_col = None if audio_columns: text_keywords = ["text", "sentence", "transcript", "transcription", "label"] @@ -489,7 +520,7 @@ def detect_multimodal_dataset(dataset): is_audio = len(audio_columns) > 0 - # Detect speaker_id column for TTS datasets (CSM, Orpheus, Spark) + # speaker_id column for TTS datasets (CSM, Orpheus, Spark) detected_speaker_col = None if audio_columns: speaker_keywords = ["source", "speaker", "speaker_id"] @@ -515,23 +546,20 @@ def _is_image_value(value) -> bool: if value is None: return False - # PIL Image instance try: from PIL.Image import Image as PILImage - if isinstance(value, PILImage): return True except ImportError: pass - # HF datasets Image feature stores decoded images as PIL or dicts with - # {"bytes": b"...", "path": "..."} when not yet decoded. + # HF Image feature: decoded as PIL, or {"bytes", "path"} when undecoded. # Exclude audio dicts (decoded audio has "array" + "sampling_rate"). if isinstance(value, dict): if "array" in value and "sampling_rate" in value: - return False # This is audio, not image + return False # audio, not image if "bytes" in value and "path" in value: - # Check path extension to exclude audio files + # Use path extension to exclude audio files. path = value.get("path") or "" if isinstance(path, str) and any( path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS @@ -539,20 +567,17 @@ def _is_image_value(value) -> bool: return False return True - # Raw bytes with a known image magic header if isinstance(value, (bytes, bytearray)): return _has_image_header(value) - # String that looks like an image file path or URL + # String that looks like an image file path or URL. _IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg") if isinstance(value, str) and len(value) < 1000: lower = value.strip().lower() - # Image URL (http://... ending in image extension) if lower.startswith(("http://", "https://")) and any( lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS ): return True - # Image file path (relative or absolute path ending in image extension) if any(lower.endswith(ext) for ext in _IMAGE_EXTS): return True @@ -577,11 +602,10 @@ def _is_audio_value(value) -> bool: if value is None: return False - # HF datasets Audio feature: decoded → {"array": np.ndarray, "sampling_rate": int} + # HF Audio feature: decoded -> {"array", "sampling_rate"}; undecoded -> {"bytes", "path"}. if isinstance(value, dict): if "array" in value and "sampling_rate" in value: return True - # Undecoded/streaming → {"bytes": b"...", "path": "some.wav"} if "bytes" in value or "path" in value: path = value.get("path") or "" if isinstance(path, str) and any( @@ -596,28 +620,22 @@ def _has_image_header(data: bytes) -> bool: """Quick magic-byte check for common image formats.""" if len(data) < 4: return False - # JPEG - if data[:2] == b"\xff\xd8": + if data[:2] == b"\xff\xd8": # JPEG return True - # PNG - if data[:4] == b"\x89PNG": + if data[:4] == b"\x89PNG": # PNG return True - # GIF - if data[:3] == b"GIF": + if data[:3] == b"GIF": # GIF return True - # WebP - if data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP": + if data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP": # WebP return True - # BMP - if data[:2] == b"BM": + if data[:2] == b"BM": # BMP return True return False def detect_vlm_dataset_structure(dataset): - """ - Detects if VLM dataset is: - - Standard VLM messages format (image objects in content) + """Detect which VLM dataset shape this is: + - Standard VLM messages (image objects in content) - Llava format (image indices + separate images column) - Simple format needing conversion (image + text columns) """ @@ -634,7 +652,6 @@ def detect_vlm_dataset_structure(dataset): column_names = set(sample.keys()) - # Check if has messages column if "messages" in column_names: messages = sample["messages"] @@ -645,11 +662,9 @@ def detect_vlm_dataset_structure(dataset): if isinstance(content, list) and len(content) > 0: if isinstance(content[0], dict) and "type" in content[0]: - # Check for llava format + # Llava format? has_index = any( - "index" in item - for item in content - if isinstance(item, dict) + "index" in item for item in content if isinstance(item, dict) ) has_images_column = "images" in column_names @@ -664,9 +679,7 @@ def detect_vlm_dataset_structure(dataset): # Standard VLM format has_image = any( - "image" in item - for item in content - if isinstance(item, dict) + "image" in item for item in content if isinstance(item, dict) ) if has_image: return { @@ -677,8 +690,8 @@ def detect_vlm_dataset_structure(dataset): "text_column": None, } - # Check for ShareGPT/ChatML conversations with placeholder + companion image column - # (e.g. Lin-Chen/ShareGPT4V, LLaVA-style datasets) + # ShareGPT/ChatML conversations with placeholder + companion + # image column (e.g. Lin-Chen/ShareGPT4V, LLaVA-style datasets) for chat_col in ("conversations", "messages"): if chat_col not in column_names: continue @@ -688,11 +701,10 @@ def detect_vlm_dataset_structure(dataset): first_msg = chat_data[0] if not isinstance(first_msg, dict): continue - # Detect ShareGPT (from/value) or ChatML (role/content) keys + # ShareGPT (from/value) or ChatML (role/content). msg_text = first_msg.get("value") or first_msg.get("content") if not isinstance(msg_text, str): continue - # Check for placeholder anywhere in the conversation has_image_placeholder = any( "" in str(m.get("value", "") or m.get("content", "")) for m in chat_data @@ -700,7 +712,7 @@ def detect_vlm_dataset_structure(dataset): ) if not has_image_placeholder: continue - # Find companion image column + # Find companion image column. image_col = None for col in column_names: if col == chat_col: @@ -717,9 +729,7 @@ def detect_vlm_dataset_structure(dataset): "messages_column": chat_col, } - # Find image and text columns using metadata filtering - - # Define metadata patterns to EXCLUDE + # Find image and text columns, filtering out metadata patterns metadata_patterns = { "suffixes": [ "_id", @@ -743,7 +753,6 @@ def detect_vlm_dataset_structure(dataset): ], } - # Image-related keywords image_keywords = [ "image", "img", @@ -756,7 +765,6 @@ def detect_vlm_dataset_structure(dataset): "filename", ] - # Text-related keywords text_keywords = [ "text", "caption", @@ -769,60 +777,48 @@ def detect_vlm_dataset_structure(dataset): ] def is_metadata_column(col_name): - """Check if column name looks like metadata.""" + """True if the column name looks like metadata.""" col_lower = col_name.lower() - # Check suffixes if any(col_lower.endswith(suffix) for suffix in metadata_patterns["suffixes"]): return True - - # Check prefixes - if any( - col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"] - ): + if any(col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"]): return True return False def _score_image_candidate(col, sample_value): """Score a candidate image column by how resolvable its value is.""" - # PIL Image object (highest priority - already loaded) + # PIL Image (already loaded) -> highest. if hasattr(sample_value, "size") and hasattr(sample_value, "mode"): return 100 - # Dict with image data (bytes/path from HF Image feature) - if isinstance(sample_value, dict) and ( - "bytes" in sample_value or "path" in sample_value - ): + # HF Image feature dict. + if isinstance(sample_value, dict) and ("bytes" in sample_value or "path" in sample_value): return 75 if isinstance(sample_value, str): - # URL strings - if sample_value.startswith(("http://", "https://")): + if sample_value.startswith(("http://", "https://")): # URL return 70 if not is_metadata_column(col) else 55 - # Bare file path - if is_metadata_column(col): + if is_metadata_column(col): # bare file path return 30 return 50 return 0 def _probe_image_candidate(col, sample_value): - """Quick probe to check if an image candidate is actually reachable. - Returns True if likely valid, False if definitely broken.""" + """Probe whether an image candidate is reachable (True unless definitely broken).""" import os - # PIL / dict — already loaded, always valid + # PIL / dict — already loaded. if not isinstance(sample_value, str): return True - # Local file — check it exists + # Local file — check it exists. if not sample_value.startswith(("http://", "https://")): - return os.path.exists( - sample_value - ) # bare filenames return False here, that's OK + return os.path.exists(sample_value) # bare filenames return False, that's OK - # URL — quick HEAD request with short timeout + # URL — quick HEAD with short timeout. try: import urllib.request @@ -833,11 +829,10 @@ def detect_vlm_dataset_structure(dataset): return False def find_image_column(): - """Find image column by keyword match + value-based fallback. - When multiple candidates exist, probes them to find one that works.""" + """Find image column by keyword match + value-based fallback, probing for one that works.""" candidates = [] - # Pass 1: keyword-matched columns + # Pass 1: keyword-matched columns. for col in column_names: if any(_keyword_in_column(keyword, col) for keyword in image_keywords): sample_value = sample[col] @@ -845,8 +840,8 @@ def detect_vlm_dataset_structure(dataset): if score > 0: candidates.append((col, score)) - # Pass 2: value-based fallback — find columns with image URLs/paths - # even if the column name doesn't match image keywords + # Pass 2: value-based fallback for image URLs/paths even when the name + # doesn't match keywords. already = {c[0] for c in candidates} for col in column_names: if col in already: @@ -854,7 +849,7 @@ def detect_vlm_dataset_structure(dataset): sample_value = sample[col] if _is_image_value(sample_value): score = _score_image_candidate(col, sample_value) - # Slightly penalise non-keyword columns so keyword matches win on ties + # Penalise non-keyword columns so keyword matches win on ties. candidates.append((col, max(score - 5, 1))) if not candidates: @@ -862,48 +857,43 @@ def detect_vlm_dataset_structure(dataset): candidates.sort(key = lambda x: x[1], reverse = True) - # Single candidate or top candidate is PIL/dict — no probing needed + # Single candidate or top is PIL/dict — no probing needed. if len(candidates) == 1 or candidates[0][1] >= 75: return candidates[0][0] - # Multiple string-based candidates — probe to find one that actually works + # Multiple string candidates — probe for one that works. for col, score in candidates: sample_value = sample[col] if _probe_image_candidate(col, sample_value): return col - # Nothing probed successfully — return highest-scored anyway and let - # conversion handle the error (it may still resolve via hf_hub_download) + # None probed OK — return highest-scored; conversion may still resolve it. return candidates[0][0] def find_text_column(): - """Find text column by filtering out metadata and checking keywords.""" + """Find text column: skip metadata, match keywords.""" candidates = [] for col in column_names: - # Skip metadata columns if is_metadata_column(col): continue - # Check if contains text keywords (word-boundary match) if any(_keyword_in_column(keyword, col) for keyword in text_keywords): - # Verify it's actually text sample_value = sample[col] if isinstance(sample_value, str) and len(sample_value) > 0: - # Longer text = higher priority (likely content, not just a label) - priority = min(len(sample_value), 1000) # Cap at 1000 + # Longer text = higher priority (content, not a label). + priority = min(len(sample_value), 1000) candidates.append((col, priority)) elif ( isinstance(sample_value, list) and len(sample_value) > 0 and isinstance(sample_value[0], str) ): - # List of strings (e.g. captions list) — lower priority than plain strings + # List of strings (e.g. captions) — lower priority than plain str. priority = min(len(sample_value[0]), 1000) // 2 candidates.append((col, priority)) - # Return highest priority candidate if candidates: candidates.sort(key = lambda x: x[1], reverse = True) return candidates[0][0] diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 4c66d2ebf6..f7b35e2869 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -1,16 +1,11 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -LLM-assisted dataset analysis using an ephemeral GGUF helper model. +"""LLM-assisted dataset analysis using an ephemeral GGUF helper model. -Complements heuristic-based detection in format_detection.py and -vlm_processing.py. Only invoked when heuristics are uncertain. - -Architecture: - - Instantiates LlamaCppBackend, loads model, runs completion(s), unloads. - - Not kept warm — VRAM is freed immediately after use. - - Gracefully degrades: returns None when unavailable (no binary, OOM, disabled). +Complements heuristic detection (format_detection.py, vlm_processing.py); only +invoked when heuristics are uncertain. Loads LlamaCppBackend, runs completion(s), +unloads (VRAM freed immediately). Degrades gracefully to None when unavailable. """ import json @@ -26,29 +21,22 @@ from loggers import get_logger logger = get_logger(__name__) -DEFAULT_HELPER_MODEL_REPO = "unsloth/gemma-4-E2B-it-GGUF" +DEFAULT_HELPER_MODEL_REPO = "unsloth/Qwen3.5-4B-MTP-GGUF" DEFAULT_HELPER_MODEL_VARIANT = "UD-Q4_K_XL" README_MAX_CHARS = 1500 def _strip_think_tags(text: str) -> str: - """Strip ... reasoning blocks emitted by some models. - - If the model places its actual answer OUTSIDE the think block, we - discard the think block and keep the rest. If the entire response - is INSIDE a think block (nothing useful outside), we extract and - return the inner content instead of discarding everything. - """ + """Strip ... blocks, keeping content outside; if all inside, return the inner.""" if "" not in text: return text - # Try stripping think blocks — keep content outside them stripped = re.sub(r".*?\s*", "", text, flags = re.DOTALL).strip() if stripped: return stripped - # Everything was inside tags — extract the inner content of the last block + # Everything was inside tags: return the last block's inner content. matches = re.findall(r"(.*?)", text, flags = re.DOTALL) if matches: return matches[-1].strip() @@ -57,20 +45,15 @@ def _strip_think_tags(text: str) -> str: def precache_helper_gguf(): - """ - Pre-download the helper GGUF to HF cache. + """Pre-download the helper GGUF to HF cache (on startup, background thread). - Called on FastAPI startup in a background thread so subsequent - ``_run_with_helper()`` calls skip the download and only pay for - llama-server startup. No-op if already cached or disabled. + Lets later ``_run_with_helper()`` calls skip the download. No-op if cached or disabled. """ if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"): return repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) - variant = os.environ.get( - "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT - ) + variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT) try: from huggingface_hub import HfApi, hf_hub_download @@ -79,16 +62,13 @@ def precache_helper_gguf(): disable_progress_bars() logging.getLogger("huggingface_hub").setLevel(logging.WARNING) - # Find the GGUF file matching the variant api = HfApi() files = api.list_repo_files(repo, repo_type = "model") gguf_files = [f for f in files if f.endswith(".gguf")] - # Find all GGUF files matching the variant (may be split into shards) + # GGUF files matching the variant (may be split into shards). variant_lower = variant.lower().replace("-", "_") - matching = sorted( - f for f in gguf_files if variant_lower in f.lower().replace("-", "_") - ) + matching = sorted(f for f in gguf_files if variant_lower in f.lower().replace("-", "_")) if matching: logger.info( @@ -110,18 +90,12 @@ def precache_helper_gguf(): def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: - """ - Load helper model, run one chat completion, unload. - - Returns the completion text, or None on any failure. - """ + """Load helper model, run one chat completion, unload. Returns text or None on failure.""" if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"): return None repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) - variant = os.environ.get( - "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT - ) + variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT) backend = None try: @@ -143,9 +117,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: return None messages = [{"role": "user", "content": prompt}] - logger.info( - "Helper model request: enable_thinking=False (per-request override)" - ) + logger.info("Helper model request: enable_thinking=False (per-request override)") cumulative = "" for chunk in backend.generate_chat_completion( messages = messages, @@ -158,7 +130,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: ): if isinstance(chunk, dict): continue # skip metadata events - cumulative = chunk # cumulative — last value is full text + cumulative = chunk # last value is full text result = cumulative.strip() result = _strip_think_tags(result) @@ -186,21 +158,10 @@ def llm_generate_vlm_instruction( samples: list[dict], dataset_name: Optional[str] = None, ) -> Optional[dict]: + """Ask a helper LLM for a task-specific VLM instruction (when heuristics are low-confidence). + + Returns {"instruction": str, "confidence": 0.85} or None. """ - Ask a helper LLM to generate a task-specific VLM instruction. - - Called when heuristic instruction generation returns low confidence - or falls back to generic. - - Args: - column_names: Column names in the dataset. - samples: 3-5 sample rows with text values (images replaced by ""). - dataset_name: Optional HF dataset identifier for context. - - Returns: - {"instruction": str, "confidence": 0.85} or None. - """ - # Format samples for the prompt formatted = "" for i, row in enumerate(samples[:5], 1): parts = [] @@ -226,9 +187,8 @@ def llm_generate_vlm_instruction( if not result: return None - # Clean up: strip quotes, ensure it's a single sentence instruction = result.strip().strip('"').strip("'").strip() - # Reject obviously bad outputs (too short, too long, or multi-line) + # Reject bad outputs (too short, too long, or multi-line). if len(instruction) < 10 or len(instruction) > 200 or "\n" in instruction: logger.warning(f"Helper model returned unusable instruction: {instruction!r}") return None @@ -240,22 +200,10 @@ def llm_generate_vlm_instruction( } -def llm_classify_columns( - column_names: list[str], - samples: list[dict], -) -> Optional[dict[str, str]]: - """ - Ask a helper LLM to classify dataset columns into roles. +def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Optional[dict[str, str]]: + """Ask a helper LLM to classify columns into roles (when heuristic detection fails). - Called when heuristic column detection fails (returns None). - - Args: - column_names: Column names in the dataset. - samples: 3-5 sample rows with values truncated to 200 chars. - - Returns: - Dict mapping column_name → role ("user"|"assistant"|"system"|"metadata"), - or None on failure. + Returns {column_name: role} for roles user|assistant|system|metadata, or None. """ formatted = "" for i, row in enumerate(samples[:5], 1): @@ -281,10 +229,9 @@ def llm_classify_columns( if not result: return None - # Parse JSON from response (may have markdown fences) + # Parse JSON from response (may have markdown fences). text = result.strip() if text.startswith("```"): - # Strip markdown code fence lines = text.split("\n") text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:]) text = text.strip() @@ -292,9 +239,7 @@ def llm_classify_columns( try: mapping = json.loads(text) except json.JSONDecodeError: - # Try to find JSON object in the response import re - match = re.search(r"\{[^}]+\}", text) if match: try: @@ -309,21 +254,17 @@ def llm_classify_columns( if not isinstance(mapping, dict): return None - # Validate: all values must be valid roles + # Keep only valid roles. valid_roles = {"user", "assistant", "system", "metadata"} cleaned = {} for col, role in mapping.items(): - if ( - col in column_names - and isinstance(role, str) - and role.lower() in valid_roles - ): + if col in column_names and isinstance(role, str) and role.lower() in valid_roles: cleaned[col] = role.lower() if not cleaned: return None - # Must have at least user + assistant + # Must have at least user + assistant. roles_present = set(cleaned.values()) if "user" not in roles_present or "assistant" not in roles_present: logger.warning(f"Helper model mapping missing user/assistant: {cleaned}") @@ -339,19 +280,9 @@ def llm_generate_dataset_warning( modality: str = "text", column_names: Optional[list[str]] = None, ) -> Optional[str]: - """ - Ask the helper LLM to turn technical dataset issues into a user-friendly warning. + """Ask the helper LLM to turn technical dataset issues into a friendly warning (any modality). - Works for all modalities (text, vision, audio). - - Args: - issues: List of technical issue descriptions found during analysis. - dataset_name: Optional HF dataset name. - modality: "text", "vision", or "audio". - column_names: Optional list of column names for context. - - Returns: - A human-friendly warning string, or None on failure. + Returns a human-friendly warning string, or None on failure. """ if not issues: return None @@ -375,7 +306,6 @@ def llm_generate_dataset_warning( return None warning = result.strip() - # Reject obviously bad outputs if len(warning) < 10 or len(warning) > 500: return None @@ -420,7 +350,11 @@ def _parse_json_response(text: str) -> Optional[dict]: return None -def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) -> str: +def _generate_with_backend( + backend, + messages: list[dict], + max_tokens: int = 512, +) -> str: """Run one chat completion on an already-loaded backend. Returns raw text.""" logger.info("Advisor request: enable_thinking=False (per-request override)") cumulative = "" @@ -431,7 +365,7 @@ def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) top_k = 20, max_tokens = max_tokens, repetition_penalty = 1.0, - enable_thinking = False, # Always disable thinking for AI Assist + enable_thinking = False, # disable thinking for AI Assist ): if isinstance(chunk, dict): continue # skip metadata events @@ -444,12 +378,7 @@ def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) def fetch_hf_dataset_card( dataset_name: str, hf_token: Optional[str] = None ) -> tuple[Optional[str], Optional[dict]]: - """ - Fetch HF dataset card (README) and metadata. - - Returns: - (readme_text, metadata_dict) or (None, None) on failure. - """ + """Fetch HF dataset card (README) and metadata. Returns (readme, metadata) or (None, None).""" try: from huggingface_hub import DatasetCard @@ -464,7 +393,7 @@ def fetch_hf_dataset_card( else: readme = readme[:README_MAX_CHARS] + "\n[...truncated]" - # Extract metadata from YAML frontmatter + # Extract metadata from YAML frontmatter. metadata = {} if card.data: for key in ( @@ -480,9 +409,7 @@ def fetch_hf_dataset_card( if val is not None: metadata[key] = val - logger.info( - f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields" - ) + logger.info(f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields") return readme, metadata except Exception as e: @@ -500,18 +427,15 @@ def _run_multi_pass_advisor( model_type: Optional[str] = None, hf_token: Optional[str] = None, ) -> Optional[dict[str, Any]]: - """ - Multi-pass LLM analysis: classify → convert → validate. + """Multi-pass LLM analysis (classify -> convert -> validate), model loaded across passes. - Keeps model loaded across all passes. Returns combined result dict or None. + Returns combined result dict or None. """ if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"): return None repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) - variant = os.environ.get( - "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT - ) + variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT) backend = None try: @@ -541,9 +465,7 @@ def _run_multi_pass_advisor( samples_text += f"Row {i}:\n" + "\n".join(parts) + "\n" metadata_str = ( - json.dumps(dataset_metadata, indent = 2, default = str)[:500] - if dataset_metadata - else "N/A" + json.dumps(dataset_metadata, indent = 2, default = str)[:500] if dataset_metadata else "N/A" ) card_excerpt = (dataset_card or "")[:1200] or "N/A" @@ -637,7 +559,7 @@ def _run_multi_pass_advisor( logger.warning(f"Advisor Pass 1 failed to produce JSON: {raw1[:200]}") return None - # If dataset is already conversational, skip passes 2-3 + # Already conversational: skip passes 2-3. if pass1.get("is_conversational") and not pass1.get("needs_conversion"): return { "success": True, @@ -742,13 +664,11 @@ def _run_multi_pass_advisor( column_roles = pass2.get("column_roles", {}) label_map = pass2.get("label_mapping") or {} # may be null - # Validate: must have at least one user AND one assistant + # Must have at least one user AND one assistant roles_present = set(column_roles.values()) if "user" not in roles_present or "assistant" not in roles_present: - logger.warning( - f"Pass 2 sanity fail: missing user or assistant role: {column_roles}" - ) - return None # triggers fallback to simple classification + logger.warning(f"Pass 2 sanity fail: missing user or assistant role: {column_roles}") + return None # falls back to simple classification # ── Pass 3: System prompt (non-conversational datasets only) ── sys_prompt = "" @@ -759,7 +679,7 @@ def _run_multi_pass_advisor( logger.info("Pass 3: Generating system prompt...") t3 = time.monotonic() - # Format label mapping info for the prompt + # Format label mapping for the prompt. label_info = "" if label_map: for col, mapping in label_map.items(): @@ -767,7 +687,6 @@ def _run_multi_pass_advisor( pairs = ", ".join(f"{k} = {v}" for k, v in mapping.items()) label_info += f"\nLabel mapping for '{col}': {pairs}" - # Describe the role assignments for context user_cols = [c for c, r in column_roles.items() if r == "user"] asst_cols = [c for c, r in column_roles.items() if r == "assistant"] task_desc = pass1.get("task_description") or pass1.get("description", "") @@ -802,18 +721,18 @@ def _run_multi_pass_advisor( ) if raw3: - # Pass 3 returns raw text, not JSON — clean it up + # Pass 3 returns raw text, not JSON. cleaned = raw3.strip().strip('"').strip("'").strip() if len(cleaned) >= 20 and cleaned.lower() not in ("null", "none", ""): sys_prompt = cleaned - # Build suggested_mapping (column → role, for the frontend dropdowns) + # Build suggested_mapping (column -> role) for the frontend dropdowns. suggested_mapping = {} for col, role in column_roles.items(): if col in columns and role in ("user", "assistant", "system"): suggested_mapping[col] = role - # Build user notification from Pass 1 classification + # Build user notification from Pass 1 classification. desc = pass1.get("task_description") or pass1.get("description", "") note_parts = [f"This is a {dtype} dataset (not conversational)."] if desc: @@ -859,23 +778,18 @@ def llm_conversion_advisor( model_name: Optional[str] = None, model_type: Optional[str] = None, ) -> Optional[dict[str, Any]]: - """ - Full conversion advisor: fetch HF card → multi-pass LLM analysis. + """Full conversion advisor: fetch HF card -> multi-pass LLM analysis. Falls back to simple llm_classify_columns() if the multi-pass advisor fails. - - Returns: - Dict with keys: success, suggested_mapping, system_prompt, user_template, - assistant_template, label_mapping, dataset_type, is_conversational, - user_notification. Or None on complete failure. + Returns a result dict (success, suggested_mapping, system_prompt, label_mapping, + dataset_type, is_conversational, user_notification, ...) or None. """ - # Fetch HF dataset card if this looks like a HF dataset (has a slash) + # Fetch HF dataset card if this looks like a HF dataset (has a slash). dataset_card = None dataset_metadata = None if dataset_name and "/" in dataset_name: dataset_card, dataset_metadata = fetch_hf_dataset_card(dataset_name, hf_token) - # Try multi-pass advisor result = _run_multi_pass_advisor( columns = column_names, samples = samples, @@ -891,7 +805,7 @@ def llm_conversion_advisor( logger.info(f"Conversion advisor succeeded: type={result.get('dataset_type')}") return result - # Fallback: simple column classification + # Fallback: simple column classification. logger.info("Advisor failed, falling back to simple column classification") simple_mapping = llm_classify_columns(column_names, samples) if simple_mapping: diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index eb2e5482c9..463d26a692 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Model and template mappings for dataset processing. +"""Model and template mappings for dataset processing. -This module contains the mapping dictionaries that associate model names -with their corresponding chat templates and response markers. +Maps model names to their chat templates and response markers. """ TEMPLATE_TO_MODEL_MAPPER = { @@ -436,7 +434,7 @@ for key, values in TEMPLATE_TO_MODEL_MAPPER.items(): for value in values: MODEL_TO_TEMPLATE_MAPPER[value] = key - # Get lowercased + # Also map lowercased names. lowered_key = key.lower() for value in values: MODEL_TO_TEMPLATE_MAPPER[value.lower()] = lowered_key @@ -445,8 +443,8 @@ for key, values in TEMPLATE_TO_MODEL_MAPPER.items(): def is_gpt_oss_model_name(name: str) -> bool: """Name-based check for gpt-oss / harmony models. - Used by both the in-process backend and the parent-process - orchestrator to detect harmony models without an IPC round-trip. + Used by the in-process backend and the parent orchestrator to detect + harmony models without an IPC round-trip. """ name = (name or "").lower() if not name: diff --git a/studio/backend/utils/datasets/raw_text.py b/studio/backend/utils/datasets/raw_text.py index 353145fd5a..03315fb287 100644 --- a/studio/backend/utils/datasets/raw_text.py +++ b/studio/backend/utils/datasets/raw_text.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Shared helpers for raw-text dataset preparation. -""" +"""Shared helpers for raw-text dataset preparation.""" from dataclasses import dataclass from typing import Literal @@ -40,10 +38,7 @@ def _split_scope(split_name: str | None) -> str: def _drop_invalid_text_rows( - dataset: Dataset, - *, - mode_title: str, - split_scope: str, + dataset: Dataset, *, mode_title: str, split_scope: str ) -> tuple[Dataset, list[RawTextNotice]]: filtered_dataset = dataset.filter(lambda ex: isinstance(ex["text"], str)) dropped_rows = len(dataset) - len(filtered_dataset) @@ -105,8 +100,7 @@ def prepare_raw_text_dataset( notices.append( RawTextNotice( message = ( - f"{mode_title}: renaming column '{renamed_col}' -> 'text' " - f"for {split_scope}" + f"{mode_title}: renaming column '{renamed_col}' -> 'text' " f"for {split_scope}" ), level = "info", ) diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py index 7b63152ede..f018913fa8 100644 --- a/studio/backend/utils/datasets/vlm_processing.py +++ b/studio/backend/utils/datasets/vlm_processing.py @@ -4,8 +4,8 @@ """ VLM (Vision-Language Model) processing utilities. -This module contains functions for generating smart instructions -for VLM datasets based on content analysis and heuristics. +Generates smart instructions for VLM datasets via content analysis and +heuristics. """ import re @@ -19,19 +19,19 @@ def generate_smart_vlm_instruction( dataset_name = None, ): """ - Generate smart, context-aware instruction for VLM datasets using heuristics. + Generate a smart, context-aware instruction for VLM datasets via heuristics. Strategy: - 1. Check for explicit question/instruction columns → use that + 1. Explicit question/instruction column → use that 2. Infer from text column name + sample content 3. Analyze dataset name for task hints - 4. Fall back to generic instruction + 4. Generic fallback Returns: dict: { "instruction": str or None, # None means use column content "instruction_type": "explicit" | "inferred" | "generic", - "uses_dynamic_instruction": bool, # True if instruction varies per sample + "uses_dynamic_instruction": bool, # True if it varies per sample "confidence": float, # 0.0 to 1.0 } """ @@ -39,16 +39,16 @@ def generate_smart_vlm_instruction( sample = next(iter(dataset)) # ===== LEVEL 1: Explicit Instruction Columns ===== - # Check for columns that contain per-sample instructions + # Columns that hold per-sample instructions question_columns = ["question", "query", "prompt", "instruction", "user_prompt"] for col in question_columns: if col in column_names: - # Check if this column has varied content (not just empty/same) + # Use it only if it has non-empty content sample_content = sample[col] if sample_content and str(sample_content).strip(): return { - "instruction": None, # Signal to use column content + "instruction": None, # use column content "instruction_column": col, "instruction_type": "explicit", "uses_dynamic_instruction": True, @@ -58,7 +58,6 @@ def generate_smart_vlm_instruction( # ===== LEVEL 2: Infer from Column Names + Content ===== text_col_lower = text_column.lower() - # Sample the text content to detect patterns text_sample = str(sample.get(text_column, ""))[:500] # First 500 chars # Task-specific keywords and their instructions @@ -66,9 +65,7 @@ def generate_smart_vlm_instruction( # OCR / Transcription "ocr": { "keywords": ["ocr", "transcribe", "transcript"], - "content_hints": [ - r"[A-Za-z\u0600-\u06FF]{10,}" - ], # Long text passages (Latin/Arabic) + "content_hints": [r"[A-Za-z\u0600-\u06FF]{10,}"], # Long Latin/Arabic passages "instruction": "Transcribe all the text shown in this image.", "confidence": 0.9, }, @@ -124,24 +121,21 @@ def generate_smart_vlm_instruction( }, } - # Check column name matches + # Score each task by column/dataset name and content matches best_match = None best_score = 0.0 for task_name, task_info in task_patterns.items(): score = 0.0 - # Check column name if any(keyword in text_col_lower for keyword in task_info["keywords"]): score += 0.5 - # Check dataset name if provided if dataset_name and any( keyword in dataset_name.lower() for keyword in task_info["keywords"] ): score += 0.3 - # Check content patterns for pattern in task_info["content_hints"]: if re.search(pattern, text_sample, re.IGNORECASE): score += 0.4 @@ -164,7 +158,6 @@ def generate_smart_vlm_instruction( if dataset_name: name_lower = dataset_name.lower() - # Common dataset name patterns if "vqa" in name_lower or "question" in name_lower: return { "instruction": "Answer the question about this image.", @@ -220,7 +213,6 @@ def generate_smart_vlm_instruction( } except Exception as e: import logging - logging.getLogger(__name__).debug(f"LLM-assisted instruction skipped: {e}") # ===== LEVEL 5: Generic Fallback ===== diff --git a/studio/backend/utils/downsample.py b/studio/backend/utils/downsample.py index bccf6a23b7..2d340ca248 100644 --- a/studio/backend/utils/downsample.py +++ b/studio/backend/utils/downsample.py @@ -12,7 +12,5 @@ def downsample(values: list[float], target_count: int) -> list[float]: return [] if target_count == 1: return [values[-1]] - indices = [ - round(i * (len(values) - 1) / (target_count - 1)) for i in range(target_count) - ] + indices = [round(i * (len(values) - 1) / (target_count - 1)) for i in range(target_count)] return [values[i] for i in indices] diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 400b5dd066..5f2b2abbcf 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Hardware detection and GPU utilities -""" +"""Hardware detection and GPU utilities.""" from . import hardware as _hardware from .hardware import ( @@ -86,8 +84,8 @@ __all__ = [ def __getattr__(name: str): - """Resolve IS_ROCM at access time so callers always see the live value - after detect_hardware() runs (it flips the flag in hardware.py).""" + """Resolve IS_ROCM lazily so callers see the live value detect_hardware() + sets in hardware.py.""" if name == "IS_ROCM": return getattr(_hardware, "IS_ROCM") raise AttributeError(name) diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index fdb1ab4520..27dfe187cc 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -3,40 +3,143 @@ """AMD GPU monitoring via amd-smi. -Mirrors the nvidia.py module structure so hardware.py can swap backends -based on IS_ROCM. All functions return the same dict shapes as their -nvidia.py counterparts. +Mirrors nvidia.py so hardware.py can swap backends based on IS_ROCM. +All functions return the same dict shapes as their nvidia.py counterparts. """ import json import math import os +import platform import re +import shutil import subprocess +import sys from typing import Any, Optional from loggers import get_logger from utils.native_path_leases import child_env_without_native_path_secret +from utils.subprocess_compat import windows_hidden_subprocess_kwargs logger = get_logger(__name__) +# amd-smi on Windows initialises the full ROCm runtime on first call, which +# can take 15-25 s on cold hardware. Linux is consistently < 2 s. +_AMD_SMI_DEFAULT_TIMEOUT = 30 if platform.system() == "Windows" else 10 -def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]: - """Run amd-smi with the given arguments and return parsed JSON, or None.""" +# Circuit breaker: stop polling amd-smi after this many consecutive failures +# (each Windows failure may pop a UAC/DiskPart elevation prompt). +_AMD_SMI_FAILURE_LIMIT = 3 +_amd_smi_consecutive_failures = 0 +_amd_smi_disabled = False + + +def _hip_sdk_present() -> bool: + """True if a HIP SDK is detectable (hipinfo on PATH or under HIP_PATH/ + ROCM_PATH), meaning amd-smi has a working runtime and runs un-elevated.""" + if shutil.which("hipinfo"): + return True + for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): + root = os.environ.get(var) + if root and os.path.exists(os.path.join(root, "bin", "hipinfo.exe")): + return True + return False + + +def _amd_smi_allowed() -> bool: + """Whether it is safe to spawn amd-smi here. + + On Windows without a working HIP runtime, amd-smi elevates a child at + runtime -- popping a UAC/DiskPart prompt that RunAsInvoker can't suppress + (its manifest is asInvoker). So only call it on Windows with a HIP SDK + present or UNSLOTH_ENABLE_AMD_SMI=1. Linux amd-smi never elevates. + """ + if platform.system() != "Windows": + return True + flag = os.environ.get("UNSLOTH_ENABLE_AMD_SMI", "").strip().lower() + if flag in ("1", "true", "yes", "on"): + return True + if flag in ("0", "false", "no", "off"): + return False + return _hip_sdk_present() + + +def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]: + """Run amd-smi with the given args and return parsed JSON, or None.""" + global _amd_smi_consecutive_failures, _amd_smi_disabled + if _amd_smi_disabled: + return None + if not _amd_smi_allowed(): + # Permanently skip amd-smi on Windows w/o a HIP SDK: every call would + # pop a UAC/DiskPart prompt (see _amd_smi_allowed). VRAM polling is then + # unavailable, but that beats the prompt. Opt back in with + # UNSLOTH_ENABLE_AMD_SMI=1. + if not _amd_smi_disabled: + logger.info( + "amd-smi disabled on Windows (no HIP SDK detected) to avoid a " + "UAC/DiskPart elevation prompt; GPU VRAM polling unavailable. " + "Set UNSLOTH_ENABLE_AMD_SMI=1 to force amd-smi." + ) + _amd_smi_disabled = True + return None + if shutil.which("amd-smi") is None: + # amd-smi does not exist on Windows (neither Adrenalin nor the HIP SDK + # ship a CLI) and can be absent on minimal Linux installs. Disable the + # poller in one step instead of burning the 3-strike circuit breaker + # on guaranteed FileNotFoundError spawns. Studio's VRAM display falls + # back to torch mem_get_info. + if not _amd_smi_disabled: + logger.info( + "amd-smi not found on PATH; GPU utilization polling via " + "amd-smi unavailable (VRAM falls back to torch mem_get_info)." + ) + _amd_smi_disabled = True + return None + _amd_env = child_env_without_native_path_secret() + if platform.system() == "Windows": + # RunAsInvoker belt-and-suspenders for any manifest-elevating helper; + # the real guard is _amd_smi_allowed() above. Mirrors install scripts. + _amd_env = {**_amd_env, "__COMPAT_LAYER": "RunAsInvoker"} try: result = subprocess.run( ["amd-smi", *args, "--json"], capture_output = True, text = True, timeout = timeout, - env = child_env_without_native_path_secret(), + env = _amd_env, + **windows_hidden_subprocess_kwargs(), ) except (OSError, subprocess.TimeoutExpired) as e: - logger.warning("amd-smi query failed: %s", e) + if isinstance(e, FileNotFoundError): + # Raced a PATH change after the which() check above; absence is + # expected on Windows (no AMD product ships an amd-smi CLI there). + logger.debug("amd-smi not found (not in PATH): %s", e) + else: + logger.warning("amd-smi query failed: %s", e) + _amd_smi_consecutive_failures += 1 + if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT: + logger.info( + "amd-smi not available (not installed; expected on HIP SDK-only systems); " + "GPU VRAM polling disabled" + ) + _amd_smi_disabled = True return None - if result.returncode != 0 or not result.stdout.strip(): + if result.returncode != 0: logger.warning("amd-smi returned code %d", result.returncode) + _amd_smi_consecutive_failures += 1 + if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT: + logger.info( + "amd-smi not available (not installed; expected on HIP SDK-only systems); " + "GPU VRAM polling disabled" + ) + _amd_smi_disabled = True return None + if not result.stdout.strip(): + # Exit 0 with no output (no GPUs visible, or a version emitting nothing + # for --json). Not a tool failure, so don't trip the circuit breaker. + logger.debug("amd-smi exited 0 but returned no output") + return None + _amd_smi_consecutive_failures = 0 # reset on success try: return json.loads(result.stdout) except json.JSONDecodeError: @@ -45,7 +148,7 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]: def _parse_numeric(value: Any) -> Optional[float]: - """Extract a numeric value from amd-smi output (may be str, int, float, or dict).""" + """Extract a numeric value from amd-smi output (str, int, float, or dict).""" if value is None: return None # Newer amd-smi versions emit {"value": 10, "unit": "W"} @@ -70,9 +173,8 @@ def _parse_memory_mb(value: Any) -> Optional[float]: """Parse a memory value from amd-smi output and return MB. Handles bare numbers (assumed MB -- the amd-smi convention on every - version we have seen), dict-shaped values with explicit units - (``{"value": 192, "unit": "GiB"}`` on newer releases), and plain - strings like ``"8192 MiB"``. + version seen), dict values with explicit units (``{"value": 192, + "unit": "GiB"}`` on newer releases), and strings like ``"8192 MiB"``. """ unit = "" raw_value = value @@ -90,8 +192,8 @@ def _parse_memory_mb(value: Any) -> Optional[float]: if num is None: return None - # Unit conversion -- GPU tools (including amd-smi) use binary units even - # when labeling them "GB" or "MB", so treat GB/GiB and MB/MiB the same. + # GPU tools use binary units even when labeled "GB"/"MB", so treat GB/GiB + # and MB/MiB the same. if "gib" in unit or "gb" in unit: return num * 1024 if "mib" in unit or "mb" in unit: @@ -102,29 +204,23 @@ def _parse_memory_mb(value: Any) -> Optional[float]: # Plain bytes return num / (1024 * 1024) - # No explicit unit -- default to MB, which is the amd-smi convention - # for bare numeric values. A previous heuristic assumed values above - # ~10M were bytes, but that misclassifies small VRAM allocations - # (e.g. 5 MB = 5,242,880 reported without a unit) as ~5 TB. Modern - # amd-smi always ships explicit units, so the heuristic branch only - # fired for legacy output where MB was already the convention. + # No explicit unit: default to MB (the amd-smi convention for bare numbers). + # A bytes-above-~10M heuristic was dropped because it misclassified small + # VRAM allocations; modern amd-smi always ships explicit units. return num def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: """Extract standardized metrics from a single GPU's amd-smi data.""" - # amd-smi metric output structure varies by version; try common paths + # Output structure varies by version; try common paths usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {})) if isinstance(usage, dict): - gpu_util = _parse_numeric( - usage.get("gfx_activity", usage.get("gpu_use_percent")) - ) + gpu_util = _parse_numeric(usage.get("gfx_activity", usage.get("gpu_use_percent"))) else: gpu_util = _parse_numeric(usage) - # Temperature -- try multiple keys in priority order. - # dict.get() returns "N/A" strings rather than falling through, - # so we must try each key and check if it parses to a real number. + # Temperature: try keys in priority order, checking each parses to a real + # number (dict.get() can return "N/A" strings rather than falling through). temp_data = gpu_data.get("temperature", {}) temp = None if isinstance(temp_data, dict): @@ -144,31 +240,24 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: power_data.get("average_socket_power", power_data.get("socket_power")), ) ) - power_limit = _parse_numeric( - power_data.get("power_cap", power_data.get("max_power_limit")) - ) + power_limit = _parse_numeric(power_data.get("power_cap", power_data.get("max_power_limit"))) else: power_draw = None power_limit = None - # VRAM -- unit-aware parsing to handle varying amd-smi output formats. - # Newer amd-smi versions may return {"value": 192, "unit": "GiB"}. - # Newer amd-smi uses "mem_usage" with "total_vram" / "used_vram" keys; - # older versions use "vram" or "fb_memory_usage" with "used" / "total". + # VRAM: unit-aware parsing across amd-smi formats. Newer versions use + # "mem_usage" with "total_vram"/"used_vram"; older use "vram" or + # "fb_memory_usage" with "used"/"total". vram_data = gpu_data.get( "mem_usage", gpu_data.get("vram", gpu_data.get("fb_memory_usage", {})), ) if isinstance(vram_data, dict): vram_used_mb = _parse_memory_mb( - vram_data.get( - "used_vram", vram_data.get("vram_used", vram_data.get("used")) - ) + vram_data.get("used_vram", vram_data.get("vram_used", vram_data.get("used"))) ) vram_total_mb = _parse_memory_mb( - vram_data.get( - "total_vram", vram_data.get("vram_total", vram_data.get("total")) - ) + vram_data.get("total_vram", vram_data.get("vram_total", vram_data.get("total"))) ) else: vram_used_mb = None @@ -176,9 +265,7 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: # Build the standardized dict (same shape as nvidia._build_gpu_metrics) vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None - vram_total_gb = ( - round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None - ) + vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None vram_util = ( round((vram_used_mb / vram_total_mb) * 100, 1) if vram_used_mb is not None and vram_total_mb is not None and vram_total_mb > 0 @@ -203,13 +290,11 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: def _has_real_metrics(metrics: dict[str, Any]) -> bool: - """Return True when ``metrics`` contains at least one non-None value. + """Return True when ``metrics`` has at least one non-None value. - ``amd-smi`` can return a zero-exit JSON envelope that is missing every - expected field (error response, unsupported card, hipless container). - In that case ``_extract_gpu_metrics`` produces a dict where every value - is ``None`` -- callers must surface this as ``available: False`` rather - than ``available: True`` with empty data. + amd-smi can return a zero-exit envelope missing every field (error, + unsupported card, hipless container), yielding an all-None dict; callers must + surface that as ``available: False``. """ return any(value is not None for value in metrics.values()) @@ -221,9 +306,8 @@ def get_physical_gpu_count() -> Optional[int]: return None if isinstance(data, list): return len(data) - # Some versions return a dict with a "gpu" / "gpus" key. Guard the - # .get() access with an isinstance check so a malformed scalar / - # string response from amd-smi cannot raise AttributeError. + # Some versions return a dict with a "gpu"/"gpus" key; guard with isinstance + # so a malformed scalar/string response can't raise AttributeError. if not isinstance(data, dict): return None gpus = data.get("gpu", data.get("gpus", [])) @@ -233,12 +317,12 @@ def get_physical_gpu_count() -> Optional[int]: def _first_visible_amd_gpu_id() -> Optional[str]: - """Return the physical AMD GPU id that should be treated as 'primary'. + """Return the physical AMD GPU id treated as 'primary'. Honours HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES - in that order (HIP respects all three). Returns ``"0"`` when none are - set, and ``None`` when the env var explicitly narrows to zero GPUs - ("" or "-1"), so callers can short-circuit to "available: False". + in that order (HIP respects all three). Returns ``"0"`` when none are set, + and ``None`` when the env var narrows to zero GPUs ("" or "-1"), so callers + can short-circuit to "available: False". """ for env_name in ( "HIP_VISIBLE_DEVICES", @@ -251,11 +335,8 @@ def _first_visible_amd_gpu_id() -> Optional[str]: raw = raw.strip() if raw == "" or raw == "-1": return None - # Filter out empty tokens after splitting. This tolerates minor - # typos like ``HIP_VISIBLE_DEVICES=",1"`` (leading comma, user - # clearly meant to narrow to device 1) while still falling - # through to the next env var when every token is empty - # (e.g. ``,,,``). + # Drop empty tokens, tolerating typos like ``",1"`` while still falling + # through to the next env var when every token is empty (``,,,``). tokens = [t.strip() for t in raw.split(",") if t.strip()] if tokens: return tokens[0] @@ -286,18 +367,15 @@ def get_primary_gpu_utilization() -> dict[str, Any]: metrics = _extract_gpu_metrics(gpu_data) if not _has_real_metrics(metrics): - # amd-smi returned a JSON envelope with no usable fields (error - # response or unsupported card). Surface as unavailable rather - # than available-with-empty-data so the UI does not render a - # ghost device. + # Envelope with no usable fields: surface as unavailable so the UI + # doesn't render a ghost device. return {"available": False} metrics["available"] = True return metrics def get_visible_gpu_utilization( - parent_visible_ids: Optional[list[int]], - parent_cuda_visible_devices: Optional[str] = None, + parent_visible_ids: Optional[list[int]], parent_cuda_visible_devices: Optional[str] = None ) -> dict[str, Any]: """Return utilization metrics for visible AMD GPUs.""" if parent_visible_ids is None: @@ -319,15 +397,11 @@ def get_visible_gpu_utilization( "index_kind": "physical", } - # Extract a device list from amd-smi's envelope. Newer versions return - # a JSON array directly, older versions return a dict with a "gpus" / - # "gpu" key wrapping the list. Guard non-dict / non-list envelopes - # (scalar / string fallbacks from malformed output) so the .get() - # access cannot raise AttributeError on an unexpected shape. + # Extract a device list across envelope shapes: a JSON array, a dict under + # "gpu_data"/"gpus"/"gpu", or a guarded scalar/string fallback. if isinstance(data, list): gpu_list = data elif isinstance(data, dict): - # Newer amd-smi wraps output in {"gpu_data": [...]} gpu_list = data.get("gpu_data", data.get("gpus", data.get("gpu", [data]))) else: gpu_list = [data] @@ -336,39 +410,36 @@ def get_visible_gpu_utilization( devices = [] for fallback_idx, gpu_data in enumerate(gpu_list): - # Skip non-dict entries defensively: if amd-smi ever ships a - # scalar inside its "gpus" array (observed on some malformed - # output), _extract_gpu_metrics would raise AttributeError on - # the first .get() call. + # Skip non-dict entries (a scalar in the array would raise AttributeError). if not isinstance(gpu_data, dict): continue - # Use AMD-reported GPU ID when available, fall back to enumeration - # index. Newer amd-smi versions wrap scalars as ``{"value": 0, - # "unit": "none"}``, so route raw_id through ``_parse_numeric`` - # which already handles bare ints, floats, strings, and that - # dict shape uniformly. - raw_id = gpu_data.get( - "gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx)) - ) + # Use the AMD-reported GPU ID, else the enumeration index. _parse_numeric + # handles bare ints/floats/strings and the {"value", "unit"} dict shape. + raw_id = gpu_data.get("gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx))) parsed_id = _parse_numeric(raw_id) if parsed_id is None: - logger.debug( - "amd-smi GPU id %r could not be parsed; falling back to " - "enumeration index %d", + logger.warning( + "amd-smi GPU id %r could not be parsed; falling back to enumeration index %d", raw_id, fallback_idx, ) idx = fallback_idx else: - idx = int(parsed_id) + rounded = round(parsed_id) + if rounded != parsed_id: + logger.warning( + "amd-smi GPU id %r parsed as non-integer %r; truncating to %d", + raw_id, + parsed_id, + rounded, + ) + idx = int(rounded) if idx not in visible_set: continue metrics = _extract_gpu_metrics(gpu_data) if not _has_real_metrics(metrics): - # Skip ghost entries: an amd-smi response that decodes to a - # dict but contains no usable fields (error envelope, etc.) - # would otherwise show up as a device row with all-None - # numbers in the UI. + # Skip ghost entries (no usable fields) so the UI doesn't show an + # all-None device row. continue metrics["index"] = idx metrics["index_kind"] = "physical" diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index ede37e2953..893d0364f7 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -16,8 +16,16 @@ Usage: ... """ +import copy +import gc +import glob import os import platform +import re +import subprocess +import sys +import types +from importlib.metadata import PackageNotFoundError, version as pkg_version import structlog from loggers import get_logger from enum import Enum @@ -31,7 +39,7 @@ logger = get_logger(__name__) class DeviceType(str, Enum): - """Supported compute backends. Inherits from str so it serializes cleanly in JSON.""" + """Supported compute backends. str subclass for clean JSON serialization.""" CUDA = "cuda" XPU = "xpu" @@ -43,22 +51,14 @@ class DeviceType(str, Enum): DEVICE: Optional[DeviceType] = None CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.) -IS_ROCM: bool = ( - False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py -) +IS_ROCM: bool = False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py def _backend_label(device: DeviceType) -> str: """Return the user-facing backend name for API responses. - Internally we still represent ROCm hosts as ``DeviceType.CUDA`` because - ROCm torch sets ``torch.cuda.is_available() = True`` and reuses the whole - ``torch.cuda.*`` API surface, so branching on ``DeviceType`` stays - consistent with the rest of the codebase. For the JSON responses served - to the Studio frontend and other clients, however, "cuda" is misleading - on an AMD machine. This helper swaps the label to ``"rocm"`` when the - module-level ``IS_ROCM`` flag is set so the UI can render the correct - backend name without every caller having to duplicate the check. + ROCm hosts stay ``DeviceType.CUDA`` internally (ROCm reuses ``torch.cuda.*``), + but "cuda" is misleading in JSON, so swap to ``"rocm"`` when ``IS_ROCM`` is set. """ if IS_ROCM and device == DeviceType.CUDA: return "rocm" @@ -69,25 +69,23 @@ def _backend_label(device: DeviceType) -> str: def is_apple_silicon() -> bool: - """Check if running on Apple Silicon hardware (pure platform check, no ML imports).""" + """True on Apple Silicon (pure platform check, no ML imports).""" return platform.system() == "Darwin" and platform.machine() == "arm64" def _has_torch() -> bool: - """Check if PyTorch is importable.""" + """True if PyTorch is importable.""" try: import torch - return True except ImportError: return False def _has_mlx() -> bool: - """Check if MLX is importable.""" + """True if MLX is importable.""" try: import mlx.core - return True except ImportError: return False @@ -95,10 +93,9 @@ def _has_mlx() -> bool: def detect_hardware() -> DeviceType: """ - Detect the best available compute device and set the module-level DEVICE global. + Detect the best compute device and set the module-level DEVICE global. - Should be called exactly once during FastAPI lifespan startup. - Safe to call multiple times (idempotent). + Call once at FastAPI lifespan startup; idempotent. Detection order: 1. CUDA (NVIDIA GPU, requires torch) @@ -112,19 +109,18 @@ def detect_hardware() -> DeviceType: # --- CUDA / ROCm: try PyTorch --- if _has_torch(): import torch - if torch.cuda.is_available(): DEVICE = DeviceType.CUDA CHAT_ONLY = False device_name = torch.cuda.get_device_properties(0).name - # Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes. - # DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP. - if getattr(torch.version, "hip", None) is not None: + # Distinguish ROCm from CUDA for display only (DeviceType stays CUDA). + # AMD SDK wheels don't set torch.version.hip, so fall back to __version__. + _hip_ver = getattr(torch.version, "hip", None) + if _hip_ver is not None or "rocm" in torch.__version__.lower(): IS_ROCM = True - print( - f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}" - ) + _hip_label = _hip_ver or torch.__version__ + print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}") else: print(f"Hardware detected: CUDA -- {device_name}") return DEVICE @@ -132,7 +128,6 @@ def detect_hardware() -> DeviceType: # --- XPU: Intel GPU --- if _has_torch(): import torch - if hasattr(torch, "xpu") and torch.xpu.is_available(): DEVICE = DeviceType.XPU CHAT_ONLY = False @@ -144,9 +139,8 @@ def detect_hardware() -> DeviceType: if is_apple_silicon() and _has_mlx(): DEVICE = DeviceType.MLX CHAT_ONLY = False - # platform.processor() runs `uname -p` which returns "i386" on most - # universal2 / Rosetta-shaped Python builds even on native arm64. - # platform.machine() is "arm64" once is_apple_silicon() has gated us. + # Use platform.machine() ("arm64"); platform.processor() returns "i386" + # on universal2 / Rosetta builds even on native arm64. chip = platform.machine() or "arm64" print(f"Hardware detected: MLX — Apple Silicon ({chip})") return DEVICE @@ -162,8 +156,8 @@ def detect_hardware() -> DeviceType: def get_device() -> DeviceType: """ - Return the detected device. Auto-detects if detect_hardware() hasn't been called yet. - Prefer calling detect_hardware() explicitly at startup instead. + Return the detected device, auto-detecting if detect_hardware() hasn't run. + Prefer calling detect_hardware() explicitly at startup. """ global DEVICE if DEVICE is None: @@ -174,10 +168,8 @@ def get_device() -> DeviceType: def clear_gpu_cache(): """ Clear GPU memory cache for the current device. - Safe to call on any platform — no-ops gracefully. + Safe on any platform — no-ops gracefully. """ - import gc - gc.collect() device = get_device() @@ -190,19 +182,17 @@ def clear_gpu_cache(): torch.cuda.ipc_collect() elif device == DeviceType.XPU: import torch - torch.xpu.synchronize() torch.xpu.empty_cache() elif device == DeviceType.MLX: - # MLX manages memory automatically; no explicit cache clear needed. - # mlx.core has no empty_cache equivalent — gc.collect() above is enough. + # MLX manages memory automatically; gc.collect() above is enough. pass def get_gpu_memory_info() -> Dict[str, Any]: """ - Get GPU memory information. - Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only environments. + Get GPU memory info. + Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only. """ device = get_device() @@ -274,16 +264,14 @@ def get_gpu_memory_info() -> Dict[str, Any]: import mlx.core as mx import psutil - # MLX uses unified memory. Total = system RAM. GPU memory used - # comes from IORegistry's AGXAccelerator (system-wide, no sudo). + # Unified memory: total = system RAM, GPU used from IORegistry AGX. total = psutil.virtual_memory().total agx = _read_apple_gpu_stats() allocated = agx.get("vram_used_bytes", 0) if agx else 0 try: info = mx.device_info() - # See detect_hardware(): platform.processor() can return "i386" - # on native arm64 Python builds, so prefer machine() as fallback. + # prefer machine(); processor() can return "i386" on native arm64. gpu_name = info.get("device_name") or platform.machine() or "arm64" except Exception: gpu_name = platform.machine() or "arm64" @@ -351,16 +339,12 @@ def get_gpu_summary() -> Dict[str, Any]: def get_package_versions() -> Dict[str, Optional[str]]: """ - Return the installed versions of key ML packages. + Return installed versions of key ML packages. - Uses importlib.metadata (stdlib) so no subprocess is needed. - CUDA version comes from torch.version.cuda. - - Returns dict with keys: unsloth, torch, transformers, cuda. - Missing packages yield None. + Uses importlib.metadata (stdlib), no subprocess. CUDA version from + torch.version.cuda. Returns dict keyed unsloth/torch/transformers/cuda; + missing packages yield None. """ - from importlib.metadata import version as pkg_version, PackageNotFoundError - packages = ("unsloth", "torch", "transformers") versions: Dict[str, Optional[str]] = {} @@ -373,7 +357,6 @@ def get_package_versions() -> Dict[str, Optional[str]]: # GPU runtime version bundled with torch try: import torch - versions["cuda"] = getattr(torch.version, "cuda", None) versions["rocm"] = getattr(torch.version, "hip", None) except Exception: @@ -417,11 +400,10 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] devices = [] for ordinal, phys_idx in enumerate(device_indices): try: - # torch uses 0-based ordinals relative to CUDA_VISIBLE_DEVICES + # torch ordinals are 0-based relative to CUDA_VISIBLE_DEVICES. props = mod.get_device_properties(ordinal) total_bytes = props.total_memory - # Prefer mem_get_info (reports system-wide usage, not just this - # process) so auto-selection accounts for other GPU consumers. + # Prefer mem_get_info (system-wide) so auto-select sees other consumers. if hasattr(mod, "mem_get_info"): free_bytes, total_bytes = mod.mem_get_info(ordinal) used_bytes = total_bytes - free_bytes @@ -445,9 +427,9 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: - """Run a query against the appropriate SMI backend (amd-smi or nvidia-smi). + """Query the appropriate SMI backend (amd-smi or nvidia-smi). - Returns the result dict if available, or None on failure/unavailability. + Returns the result dict if available, else None. """ if IS_ROCM: backend_name = "amd-smi" @@ -466,7 +448,7 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: try: func = getattr(_backend, func_name) result = func(*args, **kwargs) - if result.get("available"): + if isinstance(result, dict) and result.get("available"): return result except Exception as e: logger.warning("%s %s query failed: %s", backend_name, func_name, e) @@ -476,12 +458,9 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: def _read_apple_gpu_stats() -> Dict[str, Any]: """Query macOS IORegistry for AGX (Apple GPU) live stats. No sudo needed. - Returns dict with utilization_pct, vram_used_bytes (system-wide GPU memory). - Returns empty dict on failure. + Returns dict with utilization_pct, vram_used_bytes (system-wide GPU + memory), or empty dict on failure. """ - import subprocess - import re - try: result = subprocess.run( ["ioreg", "-r", "-c", "AGXAccelerator"], @@ -506,6 +485,133 @@ def _read_apple_gpu_stats() -> Dict[str, Any]: } +def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]: + """Query AMD GPU compute utilization via Linux DRM sysfs gpu_busy_percent.""" + if platform.system() != "Linux": + return None + try: + files = glob.glob("/sys/class/drm/card*/device/gpu_busy_percent") + if not files: + return None + values = [int(open(f).read().strip()) for f in files] + return round(sum(values) / len(values), 1) + except Exception: + return None + + +def _rocm_linux_sysfs_temp_c() -> Optional[float]: + """Query AMD GPU edge temperature via Linux DRM hwmon sysfs (temp1_input, millidegrees C).""" + if platform.system() != "Linux": + return None + try: + files = glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input") + if not files: + return None + temps = [int(open(f).read().strip()) / 1000.0 for f in files] + return round(max(temps), 1) + except Exception: + return None + + +def _rocm_linux_sysfs_power_w() -> Optional[float]: + """Query AMD GPU average power draw via Linux DRM hwmon sysfs (microwatts).""" + if platform.system() != "Linux": + return None + try: + for pattern in ( + "/sys/class/drm/card*/device/hwmon/hwmon*/power1_average", + "/sys/class/drm/card*/device/hwmon/hwmon*/power1_input", + ): + files = glob.glob(pattern) + if files: + watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files) + return round(watts, 1) + return None + except Exception: + return None + + +def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]: + """Query AMD GPU compute utilization via Windows Performance Counters (3D engine nodes).""" + if platform.system() != "Windows": + return None + try: + ps = ( + "$s=(Get-Counter '\\GPU Engine(*engtype_3D*)\\Utilization Percentage'" + " -ErrorAction SilentlyContinue).CounterSamples;" + "if($s){[math]::Min(($s|Measure-Object CookedValue -Sum).Sum,100)}else{-1}" + ) + r = subprocess.run( + ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], + capture_output = True, + text = True, + timeout = 5, + ) + if r.returncode != 0 or not r.stdout.strip(): + return None + val = float(r.stdout.strip()) + return round(val, 1) if val >= 0 else None + except Exception: + return None + + +def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]: + """Query system-wide AMD GPU VRAM via Linux DRM sysfs. + + Reads /sys/class/drm/card*/device/mem_info_vram_*, which the kernel + updates in real-time across all processes. No tools required. + Returns (used_gb, total_gb) or (None, None) on failure. + """ + if platform.system() != "Linux": + return None, None + try: + used_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_used") + total_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_total") + if not used_files or not total_files: + return None, None + used_bytes = sum(int(open(f).read().strip()) for f in used_files) + total_bytes = sum(int(open(f).read().strip()) for f in total_files) + if total_bytes == 0: + return None, None + return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2) + except Exception: + return None, None + + +def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]: + """Query system-wide dedicated GPU VRAM via Windows Performance Counters. + + Same data source as Task Manager, so cross-process usage is accurate. + Works for any GPU vendor without amd-smi or nvidia-smi. + Returns (used_gb, total_gb) or (None, None) on failure. + """ + if platform.system() != "Windows": + return None, None + try: + ps = ( + "$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'" + " -ErrorAction SilentlyContinue).CounterSamples;" + "if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}" + ) + r = subprocess.run( + ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], + capture_output = True, + text = True, + timeout = 5, + ) + if r.returncode != 0 or not r.stdout.strip(): + return None, None + used_bytes = float(r.stdout.strip()) + if used_bytes < 0: + return None, None + import torch as _torch + + total_bytes = _torch.cuda.get_device_properties(0).total_memory + return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2) + except Exception: + return None, None + + def get_gpu_utilization() -> Dict[str, Any]: """Return a live snapshot of device utilization information.""" device = get_device() @@ -514,14 +620,77 @@ def get_gpu_utilization() -> Dict[str, Any]: result = _smi_query("get_primary_gpu_utilization") if result is not None: result["backend"] = _backend_label(device) + if IS_ROCM: + # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.). + _reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec()) return result + # SMI unavailable. On Windows, use Performance Counters (Task Manager + # source) for system-wide VRAM, covering cross-process usage torch can't see. + if IS_ROCM and platform.system() == "Windows": + _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() + if _win_used is not None and _win_total is not None: + _win_util = _rocm_windows_perf_counter_gpu_util_pct() + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": _win_util, + "temperature_c": None, + "vram_used_gb": _win_used, + "vram_total_gb": _win_total, + "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) + if _win_total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + # Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed. + if IS_ROCM and platform.system() == "Linux": + _linux_used, _linux_total = _rocm_linux_sysfs_vram_gb() + if _linux_used is not None and _linux_total is not None: + _linux_util = _rocm_linux_sysfs_gpu_busy_pct() + _linux_temp = _rocm_linux_sysfs_temp_c() + _linux_power = _rocm_linux_sysfs_power_w() + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": _linux_util, + "temperature_c": _linux_temp, + "vram_used_gb": _linux_used, + "vram_total_gb": _linux_total, + "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) + if _linux_total > 0 + else None, + "power_draw_w": _linux_power, + "power_limit_w": None, + "power_utilization_pct": None, + } + # Last resort: torch mem_get_info (process-local). + _visible_spec = _get_parent_visible_gpu_spec() + _numeric_ids = _visible_spec.get("numeric_ids") or [0] + _primary_idx = [_numeric_ids[0]] if _numeric_ids else [0] + _torch_devices = _torch_get_per_device_info(_primary_idx) + if _torch_devices: + _td = _torch_devices[0] + _total = _td["total_gb"] + _used = _td["used_gb"] + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": _used, + "vram_total_gb": _total, + "vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } - # MLX path: single _read_apple_gpu_stats() call carries both VRAM-used - # bytes and GPU utilization %. psutil for unified-memory total is cheap. + # MLX: _read_apple_gpu_stats() carries both VRAM-used and GPU util%. if device == DeviceType.MLX: try: import psutil - agx = _read_apple_gpu_stats() total_bytes = psutil.virtual_memory().total except Exception as e: @@ -578,6 +747,70 @@ def get_gpu_utilization() -> Dict[str, Any]: return {"available": False, "backend": _backend_label(device)} +def _apply_unified_memory_correction( + device_metrics: Dict[str, Any], torch_info: Dict[str, Any] +) -> None: + """Per-device reconciliation: when torch reports a larger memory total + than amd-smi, overwrite the smi VRAM fields in place. + + Used by both the multi-device and primary-device reconcilers so the two + endpoints stay in sync on AMD iGPUs with unified memory. + """ + torch_total_gb = torch_info["total_gb"] + smi_total_gb = device_metrics.get("vram_total_gb") or 0.0 + if torch_total_gb > smi_total_gb: + torch_used_gb = torch_info["used_gb"] + device_metrics["vram_total_gb"] = torch_total_gb + device_metrics["vram_used_gb"] = torch_used_gb + device_metrics["vram_utilization_pct"] = ( + round((torch_used_gb / torch_total_gb) * 100, 1) if torch_total_gb > 0 else None + ) + logger.debug( + "ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with " + "torch mem_get_info total (%.2f GB) for device %s", + smi_total_gb, + torch_total_gb, + torch_info.get("index"), + ) + + +def _reconcile_rocm_unified_memory(utilization: Dict[str, Any], device_indices: list[int]) -> None: + """Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo). + + amd-smi reports only the dedicated slice; torch sees the full GTT pool. When + torch total > smi total, overwrite per-device VRAM fields with the real value. + """ + torch_devices = _torch_get_per_device_info(device_indices) + if not torch_devices: + return + torch_by_index = {td["index"]: td for td in torch_devices} + for dev in utilization.get("devices", []): + td = torch_by_index.get(dev.get("index")) + if td is None: + continue + _apply_unified_memory_correction(dev, td) + + +def _reconcile_primary_rocm_unified_memory( + utilization: Dict[str, Any], parent_visible_spec: Dict[str, Any] +) -> None: + """Same fix as _reconcile_rocm_unified_memory for the flat primary-GPU dict.""" + numeric_ids = parent_visible_spec.get("numeric_ids") + if numeric_ids is None: + # No visibility env var set: torch ordinal 0 is the primary device. + primary_idx = [0] + elif len(numeric_ids) == 0: + # Empty mask: no GPU visible. Querying torch device 0 would raise or + # return stale data, so bail rather than write bad values. + return + else: + primary_idx = [int(numeric_ids[0])] + torch_devices = _torch_get_per_device_info(primary_idx) + if not torch_devices: + return + _apply_unified_memory_correction(utilization, torch_devices[0]) + + def get_visible_gpu_utilization() -> Dict[str, Any]: device = get_device() @@ -590,13 +823,16 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: ) if result is not None: result["backend"] = _backend_label(device) + numeric_ids = parent_visible_spec.get("numeric_ids") + if IS_ROCM and numeric_ids is not None: + # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.). + _reconcile_rocm_unified_memory(result, numeric_ids) return result # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) if device in (DeviceType.CUDA, DeviceType.XPU): parent_ids = get_parent_visible_gpu_ids() - # When parent_visible_ids is empty (UUID/MIG mask or no CVD set), - # enumerate torch-visible ordinals so the UI still shows devices. + # Empty parent_ids (UUID/MIG mask or no CVD): enumerate torch ordinals. if parent_ids: torch_indices = parent_ids index_kind = "physical" @@ -683,13 +919,16 @@ _visible_gpu_count: Optional[int] = None def _get_parent_visible_gpu_spec() -> Dict[str, Any]: - # ROCm uses HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES in addition to - # CUDA_VISIBLE_DEVICES (which HIP also respects). Check ROCm-specific - # env vars first so multi-GPU AMD setups are handled correctly. - # Use explicit None checks (not `or`) so empty string "" is honoured - # as "no visible GPUs" rather than falling through to CUDA_VISIBLE_DEVICES. + # ROCm uses HIP/ROCR_VISIBLE_DEVICES on top of CUDA_VISIBLE_DEVICES; check + # them first. Explicit None checks (not `or`) so "" reads as "no visible GPUs". cuda_visible = None - if IS_ROCM: + # Prefer ROCm masks only on a ROCm host or when no CUDA mask is set, so a + # stale HIP_VISIBLE_DEVICES on NVIDIA can't override CUDA_VISIBLE_DEVICES. + _is_rocm_spec = IS_ROCM or ( + "CUDA_VISIBLE_DEVICES" not in os.environ + and ("HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ) + ) + if _is_rocm_spec: hip_vis = os.environ.get("HIP_VISIBLE_DEVICES") rocr_vis = os.environ.get("ROCR_VISIBLE_DEVICES") if hip_vis is not None: @@ -762,7 +1001,7 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: f"Parent-visible GPUs: {parent_visible_ids}" ) - # Reject negative IDs unconditionally. + # Reject negative IDs. negative_ids = [gpu_id for gpu_id in requested_ids if gpu_id < 0] if negative_ids: raise ValueError( @@ -770,19 +1009,14 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: f"Rejected IDs: {negative_ids}. Parent-visible GPUs: {parent_visible_ids}" ) - # Only enforce the physical upper bound when we have a reliable count - # from nvidia-smi. When the count comes from torch, it reflects visible - # devices (filtered by CUDA_VISIBLE_DEVICES), not the physical total, - # so high physical indices like 3 would be falsely rejected on a - # CUDA_VISIBLE_DEVICES="2,3" machine that reports device_count()=2. - # The parent-visible check below is authoritative in all cases. + # Only enforce the physical upper bound when the count is reliable (nvidia-smi). + # A torch count reflects only visible devices, so it could falsely reject valid + # physical indices. The parent-visible check below is always authoritative. if physical_gpu_count > 0 and parent_visible_ids: max_parent_id = max(parent_visible_ids) if physical_gpu_count > max_parent_id: - # Count is plausibly physical (not just visible), so enforce it - out_of_range = [ - gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count - ] + # Count is plausibly physical, so enforce it. + out_of_range = [gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count] if out_of_range: raise ValueError( f"Invalid gpu_ids {requested_ids}: IDs must be physical GPU IDs " @@ -790,9 +1024,7 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: f"Rejected IDs: {out_of_range}. Parent-visible GPUs: {parent_visible_ids}" ) - disallowed_ids = [ - gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids - ] + disallowed_ids = [gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids] if disallowed_ids: raise ValueError( f"Invalid gpu_ids {requested_ids}: requested GPUs {disallowed_ids} are " @@ -813,9 +1045,7 @@ def _resolve_model_identifier_for_gpu_estimate( return config.base_model return config.identifier if config else model_name except Exception as e: - logger.debug( - "Could not resolve base model for GPU estimate '%s': %s", model_name, e - ) + logger.debug("Could not resolve base model for GPU estimate '%s': %s", model_name, e) return model_name @@ -852,7 +1082,6 @@ def _get_hf_safetensors_total_params( def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None): try: from transformers import AutoConfig - trust_remote_code = model_name.lower().startswith("unsloth/") return AutoConfig.from_pretrained( model_name, @@ -865,17 +1094,61 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non def _determine_attention_impl_for_gpu_estimate(config) -> str: - import copy as _copy + # torch.distributed is incomplete on Windows ROCm (torch._C._distributed_c10d + # can't be imported). Inject stubs into sys.modules before importing + # torch.distributed, then patch the missing process-group helpers. + if sys.platform == "win32" and IS_ROCM: + # Dummy for any name torch.distributed imports from these stubs. + class _Dummy: + pass + + for _c10d_name in ( + "torch._C._distributed_c10d", + "torch._C._distributed_autograd", + "torch._C._distributed_rpc", + ): + if _c10d_name not in sys.modules: + _stub = types.ModuleType(_c10d_name) + # No-op dummies for names torch.distributed imports from _distributed_c10d. + for _sym in ( + "FakeProcessGroup", + "ProcessGroup", + "Work", + "Store", + "PrefixStore", + "FileStore", + "TCPStore", + "HashStore", + "Reducer", + "Logger", + "DistributedDebugLevel", + "GradBucket", + "BuiltinCommHookType", + ): + setattr(_stub, _sym, _Dummy) + sys.modules[_c10d_name] = _stub + + try: + import torch.distributed as _td + for _attr, _stub in ( + ("is_initialized", lambda: False), + ("is_available", lambda: False), + ("get_rank", lambda: 0), + ("get_world_size", lambda: 1), + ("is_torchelastic_launched", lambda: False), + ): + if not hasattr(_td, _attr): + setattr(_td, _attr, _stub) + except ImportError: + pass from unsloth.models._utils import resolve_attention_implementation from transformers import AutoModel, AutoModelForCausalLM - # why: resolve_attention_implementation calls _set_attn_impl which writes - # _attn_implementation onto the config; PreTrainedConfig's setter walks - # `sub_configs` and propagates to nested text_config / sub-configs, so a - # shallow copy still mutates those shared inner objects on the cached - # config returned by _load_config_for_gpu_estimate. Deepcopy isolates them. - config_copy = _copy.deepcopy(config) + # why: resolve_attention_implementation writes _attn_implementation onto the + # config and propagates to nested sub-configs; a shallow copy would still + # mutate the cached config's shared inner objects. Deepcopy isolates them. + config_copy = copy.deepcopy(config) model_class = None for auto_model in (AutoModelForCausalLM, AutoModel): @@ -917,17 +1190,15 @@ def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]: synthetic_total_bytes, synthetic_total_bytes, ) - _, _, _, memory_left_for_kv_cache_gb = ( - _vllm_utils.approximate_vllm_memory_usage( - config, - load_in_4bit = False, - load_in_8bit = False, - max_seq_length = 1, - gpu_memory_utilization = 1.0, - enable_lora = False, - account_for_gradients = False, - cuda_graph_overhead = False, - ) + _, _, _, memory_left_for_kv_cache_gb = _vllm_utils.approximate_vllm_memory_usage( + config, + load_in_4bit = False, + load_in_8bit = False, + max_seq_length = 1, + gpu_memory_utilization = 1.0, + enable_lora = False, + account_for_gradients = False, + cuda_graph_overhead = False, ) finally: _vllm_utils.get_mem_info = original_get_mem_info @@ -949,15 +1220,11 @@ def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]: def estimate_fp16_model_size_bytes( model_name: str, hf_token: Optional[str] = None ) -> tuple[Optional[int], str]: - estimate_model = _resolve_model_identifier_for_gpu_estimate( - model_name, hf_token = hf_token - ) + estimate_model = _resolve_model_identifier_for_gpu_estimate(model_name, hf_token = hf_token) total_params = None if "/" in estimate_model and not Path(estimate_model).exists(): - total_params = _get_hf_safetensors_total_params( - estimate_model, hf_token = hf_token - ) + total_params = _get_hf_safetensors_total_params(estimate_model, hf_token = hf_token) if total_params: return int(total_params * 2), "safetensors" @@ -1012,9 +1279,7 @@ def estimate_required_model_memory_gb( DEFAULT_TARGET_MODULES, ) - model_size_bytes, source = estimate_fp16_model_size_bytes( - model_name, hf_token = hf_token - ) + model_size_bytes, source = estimate_fp16_model_size_bytes(model_name, hf_token = hf_token) metadata: Dict[str, Any] = { "mode": "inference" if training_type is None else "training", "model_size_source": source, @@ -1037,9 +1302,7 @@ def estimate_required_model_memory_gb( return required_gb, metadata training_method = ( - "full" - if training_type == "Full Finetuning" - else ("qlora" if load_in_4bit else "lora") + "full" if training_type == "Full Finetuning" else ("qlora" if load_in_4bit else "lora") ) vram_config = TrainingVramConfig( training_method = training_method, @@ -1052,38 +1315,37 @@ def estimate_required_model_memory_gb( load_in_4bit = load_in_4bit, ) - estimate_model = _resolve_model_identifier_for_gpu_estimate( - model_name, hf_token = hf_token - ) + estimate_model = _resolve_model_identifier_for_gpu_estimate(model_name, hf_token = hf_token) config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token) if config is not None: try: - vram_config.attention_implementation = ( - _determine_attention_impl_for_gpu_estimate(config) + vram_config.attention_implementation = _determine_attention_impl_for_gpu_estimate( + config ) except Exception as e: - logger.warning( + # Debug-level: fires every estimate on Windows ROCm (stub lacks Store); + # expected and non-actionable -- eager is the safe fallback. + logger.debug( "Could not resolve attention implementation for '%s': %s", estimate_model, e, ) - # why: if we cannot prove flash attention is usable, charge the - # quadratic non-flash activation path so GPU selection stays - # conservative. + # why: charge the quadratic non-flash activation path so GPU + # selection stays conservative when flash attn isn't proven usable. vram_config.attention_implementation = "eager" arch = extract_arch_config(config) if config is not None else None if arch is not None: breakdown = estimate_training_vram(arch, vram_config) - # why: extract_arch_config only sees text_config; safetensors include - # vision/audio tower bytes that the text-arch fp16 total misses. + # why: extract_arch_config only sees text_config; add the vision/audio + # tower bytes that the text-arch fp16 total misses. arch_fp16_bytes = compute_total_params(arch) * 2 extra_bytes = max(0, int(model_size_bytes) - arch_fp16_bytes) if extra_bytes > 0: breakdown.model_weights += extra_bytes if training_method == "full": - # why: full fine-tuning makes the extra (vision/audio) params - # trainable; optimizer + gradient bytes scale with them too. + # why: full fine-tuning makes extra params trainable; optimizer + + # gradient bytes scale with them. extra_params = extra_bytes // 2 breakdown.optimizer_states += compute_optimizer_bytes( extra_params, @@ -1102,7 +1364,7 @@ def estimate_required_model_memory_gb( ) return required_gb, metadata - # Fallback when model config is unavailable + # Fallback when model config is unavailable. overhead_gb = CUDA_OVERHEAD_BYTES / (1024**3) if training_method == "full": required_gb = model_size_gb * 3.5 + overhead_gb @@ -1162,9 +1424,7 @@ def auto_select_gpu_ids( return None, metadata if required_gb is None: - # Cannot estimate model size -- fall back to all visible GPUs - # rather than risk loading on a single GPU that may not have - # enough memory. + # Can't estimate size -- use all visible GPUs rather than risk one too small. parent_ids = get_parent_visible_gpu_ids() metadata["selection_mode"] = "fallback_all" metadata["selected_gpu_ids"] = parent_ids @@ -1202,17 +1462,13 @@ def auto_select_gpu_ids( free_by_index = {item["index"]: item["free_gb"] for item in ranked} selected: list[int] = [] usable_gb = 0.0 - # Multi-GPU sharding has overhead from inter-GPU communication (NCCL - # all-reduce, PCIe/NVLink transfers, synchronization barriers), so each - # additional GPU contributes less than its raw free memory. The first GPU - # keeps its full capacity (no cross-device overhead). 0.85 was calibrated - # empirically on 2-8 GPU setups with NVLink and PCIe topologies -- the - # 15% discount accounts for NCCL buffers (~2-5% of VRAM), pipeline bubble - # overhead, and memory fragmentation from non-uniform shard sizes. + # Sharding has inter-GPU overhead, so each extra GPU contributes less than + # its raw free memory (first GPU keeps full capacity). 0.85 is empirical on + # 2-8 GPU setups: covers NCCL buffers, pipeline bubbles, fragmentation. multi_gpu_overhead = 0.85 - # Per-GPU check: activations don't shard, so each GPU needs its weight - # shard + full activation cost. Use precomputed min_per_gpu_N values. + # Per-GPU check: activations don't shard, so each GPU needs its weight shard + # + full activation cost. Uses precomputed min_per_gpu_N values. vram_breakdown = estimate_metadata.get("vram_breakdown", {}) for candidate in ranked: @@ -1249,10 +1505,8 @@ def auto_select_gpu_ids( ) return selected, metadata - # Use only GPUs with verified VRAM data (from gpu_candidates, not raw devices) - fallback_all = ( - [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids - ) + # Use only GPUs with verified VRAM data. + fallback_all = [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids metadata["selection_mode"] = "fallback_all" if ranked: fallback_usable = ranked[0]["free_gb"] + sum( @@ -1290,18 +1544,17 @@ def prepare_gpu_selection( """Resolve which physical GPUs to use for a model load. GPU selection modes: - - **Explicit** (``gpu_ids=[5, 6, 7]``): the caller chooses exact GPUs. - All listed GPUs are used and the model is sharded across them via - ``device_map="balanced"``, regardless of whether the model would fit - on fewer GPUs. IDs are validated against the parent-visible set. - - **Auto** (``gpu_ids=None`` or ``[]``): ``auto_select_gpu_ids`` estimates - VRAM requirements and picks the *minimum* number of GPUs needed, - preferring GPUs with the most free memory. + - **Explicit** (``gpu_ids=[5, 6, 7]``): caller chooses exact GPUs. + All listed GPUs are used and the model is sharded via + ``device_map="balanced"``, even if it would fit on fewer. IDs are + validated against the parent-visible set. + - **Auto** (``gpu_ids=None`` or ``[]``): ``auto_select_gpu_ids`` + estimates VRAM needs and picks the *minimum* GPUs needed, + preferring those with the most free memory. - The returned ``gpu_ids`` list is later passed to ``get_device_map()`` which - maps it to a Hugging Face ``device_map`` string, and to ``apply_gpu_ids()`` - in the worker subprocess which narrows ``CUDA_VISIBLE_DEVICES`` before any - torch/CUDA initialisation. + The returned ``gpu_ids`` is later passed to ``get_device_map()`` (maps it + to a Hugging Face ``device_map`` string) and to ``apply_gpu_ids()`` in the + worker subprocess (narrows ``CUDA_VISIBLE_DEVICES`` before torch/CUDA init). """ if gpu_ids and get_device() != DeviceType.CUDA: raise ValueError( @@ -1337,8 +1590,7 @@ def get_physical_gpu_count() -> int: Return the number of physical GPUs on the machine. Uses ``nvidia-smi -L`` on NVIDIA (unaffected by CUDA_VISIBLE_DEVICES), - with a torch-based fallback for AMD ROCm and Intel XPU. - Result is cached after the first call. + with a torch fallback for AMD ROCm and Intel XPU. Cached after first call. """ global _physical_gpu_count if _physical_gpu_count is not None: @@ -1358,7 +1610,7 @@ def get_physical_gpu_count() -> int: return _physical_gpu_count except Exception: pass - # SMI tool unavailable or failed -- fall back to torch + # SMI unavailable -- fall back to torch. count = _torch_get_physical_gpu_count() _physical_gpu_count = count if count is not None else 1 return _physical_gpu_count @@ -1380,10 +1632,10 @@ def get_physical_gpu_count() -> int: def _backend_visible_devices_env() -> Optional[str]: """Return the raw visibility env string that applies to this backend. - On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence - over CUDA_VISIBLE_DEVICES; the helper mirrors the resolution logic in - ``_get_parent_visible_gpu_spec`` so ``backend_cuda_visible_devices`` - reports the value that is actually narrowing the visible device set. + On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over + CUDA_VISIBLE_DEVICES; this mirrors ``_get_parent_visible_gpu_spec`` so + ``backend_cuda_visible_devices`` reports the value actually narrowing the + visible device set. """ if IS_ROCM: return _get_parent_visible_gpu_spec().get("raw") @@ -1394,7 +1646,7 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: device = get_device() if device in (DeviceType.CUDA, DeviceType.XPU): parent_visible_ids = get_parent_visible_gpu_ids() - # Try native SMI tool first (nvidia-smi for NVIDIA, skipped for ROCm) + # Try native SMI first (nvidia-smi; skipped for ROCm). if device == DeviceType.CUDA and not IS_ROCM: try: from . import nvidia @@ -1410,9 +1662,8 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: except Exception as e: logger.warning("Backend GPU visibility query failed: %s", e) - # Torch fallback (AMD ROCm, Intel XPU, nvidia-smi missing/failed) - # When parent_visible_ids is empty (UUID/MIG mask), enumerate by - # torch ordinal so the UI still shows devices. + # Torch fallback (ROCm, XPU, nvidia-smi missing). Empty parent_visible_ids + # (UUID/MIG mask) -> enumerate by torch ordinal so the UI shows devices. if parent_visible_ids: torch_indices = parent_visible_ids index_kind = "physical" @@ -1493,15 +1744,15 @@ def get_visible_gpu_count() -> int: Return the number of GPUs visible to this process. Respects ``CUDA_VISIBLE_DEVICES`` -- if set, only those GPUs count. - Falls back to physical count if the env var is unset or torch is - unavailable. Result is cached after the first call. + Falls back to physical count if unset or torch is unavailable. + Cached after the first call. """ global _visible_gpu_count if _visible_gpu_count is not None: return _visible_gpu_count - # Use _get_parent_visible_gpu_spec() which already handles - # HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES on ROCm. + # _get_parent_visible_gpu_spec() already handles HIP_VISIBLE_DEVICES / + # ROCR_VISIBLE_DEVICES on ROCm. visible_spec = _get_parent_visible_gpu_spec() if visible_spec["raw"] is not None: raw = visible_spec["raw"].strip() @@ -1513,10 +1764,9 @@ def get_visible_gpu_count() -> int: _visible_gpu_count = len([x for x in raw.split(",") if x.strip()]) return _visible_gpu_count - # No visibility env var set -- try torch, fall back to physical count + # No visibility env var set -- try torch, else physical count try: import torch - if get_device() == DeviceType.XPU and hasattr(torch, "xpu"): _visible_gpu_count = torch.xpu.device_count() else: @@ -1531,8 +1781,7 @@ def apply_gpu_ids(gpu_ids) -> None: if gpu_ids is None: return - # Empty list means "no GPUs visible" -- treat the same as None - # (inherit parent) to avoid setting CUDA_VISIBLE_DEVICES="" which + # Empty list -> treat like None (inherit parent); setting CUDA_VISIBLE_DEVICES="" # disables CUDA entirely and crashes downstream torch calls. if isinstance(gpu_ids, (list, tuple)) and len(gpu_ids) == 0: return @@ -1545,56 +1794,68 @@ def apply_gpu_ids(gpu_ids) -> None: value = str(gpu_ids) os.environ["CUDA_VISIBLE_DEVICES"] = value - # Keep ROCm visibility env vars in sync so _get_parent_visible_gpu_spec() - # picks up the narrowed set on AMD systems. Workers can call - # apply_gpu_ids() before detect_hardware() runs (so IS_ROCM is still - # its default False), so also mirror the selection whenever the - # parent process already set a ROCm visibility variable -- that - # way a downstream ROCm process inherits the narrowed mask even - # before Studio's hardware detection has classified the host. + # Keep ROCm visibility env vars in sync. Workers may call apply_gpu_ids() + # before detect_hardware() (IS_ROCM still False), so also mirror when the + # parent set a ROCm visibility var, with a torch.version.hip probe fallback. _inherits_rocm_visibility = ( "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ ) - if IS_ROCM or _inherits_rocm_visibility: + _is_rocm = IS_ROCM or _inherits_rocm_visibility + if not _is_rocm: + # torch.version.hip is set on ROCm, None on CUDA; AMD SDK wheels may leave + # it unset but encode "rocm" in __version__. Broad except: never crash a worker. + try: + import torch as _torch + _is_rocm = ( + getattr(_torch.version, "hip", None) is not None + or "rocm" in getattr(_torch, "__version__", "").lower() + ) + except Exception as e: + logger.debug( + "apply_gpu_ids: torch ROCm probe skipped (%s: %s)", + type(e).__name__, + e, + ) + if _is_rocm: os.environ["HIP_VISIBLE_DEVICES"] = value - os.environ["ROCR_VISIBLE_DEVICES"] = value + # ROCR_VISIBLE_DEVICES operates at the HSA agent level and uses + # different indexing semantics to HIP_VISIBLE_DEVICES. Setting it + # to a physical GPU index breaks multi-GPU ROCm systems where the + # parent already set ROCR_VISIBLE_DEVICES (e.g. "0,1"): narrowing + # to "1" causes torch.cuda.is_available() to return False in the + # worker subprocess. HIP_VISIBLE_DEVICES is sufficient for GPU + # selection on ROCm -- leave ROCR_VISIBLE_DEVICES inherited. _visible_gpu_count = None - if IS_ROCM or _inherits_rocm_visibility: + if _is_rocm: logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value) else: logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value) -def get_device_map( - gpu_ids: Optional[list[int]] = None, -) -> str: +def get_device_map(gpu_ids: Optional[list[int]] = None) -> str: """Return the Hugging Face ``device_map`` string for model loading. Returns ``"balanced"`` (shard evenly across GPUs) when: - ``gpu_ids`` explicitly lists >1 GPU, **or** - ``CUDA_VISIBLE_DEVICES`` uses UUID/MIG identifiers (non-numeric) and - more than one GPU is visible (fallback: we cannot resolve numeric IDs, - so we assume the caller intends multi-GPU). + >1 GPU is visible (fallback: numeric IDs unresolvable, so assume + multi-GPU is intended). - Returns ``"sequential"`` (single device) in all other cases, including - non-CUDA backends (CPU, MLX). + Returns ``"sequential"`` (single device) otherwise, including non-CUDA + backends (CPU, MLX). - Callers should use ``prepare_gpu_selection()`` upstream to determine the - ``gpu_ids`` list -- that function handles the smart auto-selection of the - minimum number of GPUs needed for a given model. + Use ``prepare_gpu_selection()`` upstream to determine ``gpu_ids`` -- it + handles auto-selecting the minimum GPUs needed for a model. """ device = get_device() if device == DeviceType.CUDA: multi_gpu = gpu_ids is not None and len(gpu_ids) > 1 if not multi_gpu: - # UUID/MIG masks cannot be split into numeric IDs, so if multiple - # GPUs are visible we assume multi-GPU sharding is intended. + # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU + # means multi-GPU sharding is intended. parent_visible_spec = _get_parent_visible_gpu_spec() - if ( - parent_visible_spec["numeric_ids"] is None - and get_visible_gpu_count() > 1 - ): + if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1: multi_gpu = True if multi_gpu: @@ -1614,14 +1875,16 @@ def get_offloaded_device_map_entries(model) -> dict[str, str]: } -def raise_if_offloaded(model, device_map: str, context: str = "Loading") -> None: +def raise_if_offloaded( + model, + device_map: str, + context: str = "Loading", +) -> None: """Raise ``ValueError`` if *model* has modules offloaded to CPU or disk.""" offloaded = get_offloaded_device_map_entries(model) if not offloaded: return - example = ", ".join( - f"{name}={placement}" for name, placement in list(offloaded.items())[:5] - ) + example = ", ".join(f"{name}={placement}" for name, placement in list(offloaded.items())[:5]) raise ValueError( f"{context} does not support models loaded with CPU or disk offload. " f"device_map='{device_map}' produced offloaded modules: {example}" @@ -1632,18 +1895,14 @@ def safe_num_proc(desired: Optional[int] = None) -> int: """ Return a safe ``num_proc`` for ``dataset.map()`` calls. - On Windows, always returns 1 because Python uses ``spawn`` instead of - ``fork`` for multiprocessing -- the overhead of re-importing torch, - transformers, unsloth etc. per worker is typically slower than - single-process for normal dataset sizes. + On Windows always returns 1: Python uses ``spawn`` not ``fork``, so + re-importing torch/transformers/unsloth per worker is typically slower + than single-process for normal dataset sizes. - On multi-GPU machines (where multiple GPUs are *visible* to this - process) the NVIDIA driver spawns extra background threads, making - ``os.fork()`` prone to deadlocks when many workers are created. - This helper caps ``num_proc`` to 4 on such machines. - - When ``CUDA_VISIBLE_DEVICES`` restricts to a single GPU, the cap - does not apply. + On multi-GPU machines (multiple GPUs *visible* to this process) the + NVIDIA driver spawns extra background threads, making ``os.fork()`` + deadlock-prone with many workers, so this caps ``num_proc`` to 4. + The cap does not apply when ``CUDA_VISIBLE_DEVICES`` restricts to one GPU. Args: desired: The num_proc you *want*. If None, auto-computes from @@ -1652,11 +1911,8 @@ def safe_num_proc(desired: Optional[int] = None) -> int: Returns: A safe integer ≥ 1. """ - import sys - - # Windows and macOS use 'spawn' for multiprocessing -- the overhead of - # re-importing torch/transformers/unsloth per worker is typically slower - # than single-process. + # Windows/macOS use 'spawn'; re-importing torch/transformers/unsloth per + # worker is typically slower than single-process. if sys.platform in ("win32", "darwin"): return 1 @@ -1679,9 +1935,8 @@ def safe_thread_num_proc(desired: Optional[int] = None) -> int: """ Return a safe worker count for ``ThreadPoolExecutor`` calls. - Unlike ``safe_num_proc()``, this does NOT cap to 1 on macOS/Windows. - Threads share the parent process address space and are unaffected by - the ``spawn`` vs ``fork`` distinction. + Unlike ``safe_num_proc()``, does NOT cap to 1 on macOS/Windows: threads + share the parent address space, unaffected by ``spawn`` vs ``fork``. Args: desired: The thread count you *want*. If None, auto-computes @@ -1700,12 +1955,10 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]: """ Return a safe ``num_proc`` for ``Dataset.map()`` and ``Dataset.filter()``. - Returns ``None`` on spawn-based platforms (Windows, macOS) because - ``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``). - Only ``num_proc=None`` guarantees in-process execution. + Returns ``None`` on spawn platforms (Windows, macOS) because ``datasets`` + treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``); only + ``num_proc=None`` guarantees in-process execution. """ - import sys - if sys.platform in ("win32", "darwin"): return None return safe_num_proc(desired) diff --git a/studio/backend/utils/hardware/nvidia.py b/studio/backend/utils/hardware/nvidia.py index 099c5fa3a5..f98ca4343e 100644 --- a/studio/backend/utils/hardware/nvidia.py +++ b/studio/backend/utils/hardware/nvidia.py @@ -25,20 +25,12 @@ def _parse_smi_value(raw: str): def _build_gpu_metrics( - vram_used_mb, - vram_total_mb, - power_draw, - power_limit, - **extra, + vram_used_mb, vram_total_mb, power_draw, power_limit, **extra ) -> dict[str, Any]: return { **extra, - "vram_used_gb": round(vram_used_mb / 1024, 2) - if vram_used_mb is not None - else None, - "vram_total_gb": round(vram_total_mb / 1024, 2) - if vram_total_mb is not None - else None, + "vram_used_gb": round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None, + "vram_total_gb": round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None, "vram_utilization_pct": round((vram_used_mb / vram_total_mb) * 100, 1) if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0 else None, @@ -50,9 +42,7 @@ def _build_gpu_metrics( } -def _visible_ordinal_map( - parent_visible_ids: Optional[list[int]], -) -> Optional[dict[int, int]]: +def _visible_ordinal_map(parent_visible_ids: Optional[list[int]]) -> Optional[dict[int, int]]: if parent_visible_ids is None: return None return {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)} @@ -118,12 +108,10 @@ def get_primary_gpu_utilization() -> dict[str, Any]: def get_visible_gpu_utilization( - parent_visible_ids: Optional[list[int]], - parent_cuda_visible_devices: Optional[str] = None, + parent_visible_ids: Optional[list[int]], parent_cuda_visible_devices: Optional[str] = None ) -> dict[str, Any]: - # When parent_visible_ids is None (UUID/MIG mask), we cannot safely - # map nvidia-smi rows to the process's visible devices. Return empty - # instead of exposing all physical GPUs. + # parent_visible_ids None (UUID/MIG mask): can't map nvidia-smi rows to + # visible devices, so return empty rather than exposing all physical GPUs. if parent_visible_ids is None: return { "available": False, @@ -188,9 +176,7 @@ def get_visible_gpu_utilization( index = idx, index_kind = "physical", visible_ordinal = ( - visible_ordinals[idx] - if visible_ordinals is not None - else len(devices) + visible_ordinals[idx] if visible_ordinals is not None else len(devices) ), gpu_utilization_pct = _parse_smi_value(parts[1]), temperature_c = _parse_smi_value(parts[2]), @@ -207,11 +193,10 @@ def get_visible_gpu_utilization( def get_backend_visible_gpu_info( - parent_visible_ids: Optional[list[int]], - backend_cuda_visible_devices: Optional[str], + parent_visible_ids: Optional[list[int]], backend_cuda_visible_devices: Optional[str] ) -> dict[str, Any]: - # When parent_visible_ids is None (UUID/MIG mask), we cannot safely - # map nvidia-smi rows to the process's visible devices. + # parent_visible_ids None (UUID/MIG mask): can't map nvidia-smi rows to + # visible devices. if parent_visible_ids is None: return { "available": False, @@ -263,7 +248,7 @@ def get_backend_visible_gpu_info( continue if visible_ordinals is not None and idx not in visible_ordinals: continue - # Use split with limit to handle GPU names containing commas + # Rejoin in case the GPU name contains commas name = parts[1] if len(parts) == 3 else ", ".join(parts[1:-1]) try: mem_total_mb = int(parts[-1]) @@ -274,9 +259,7 @@ def get_backend_visible_gpu_info( "index": idx, "index_kind": "physical", "visible_ordinal": ( - visible_ordinals[idx] - if visible_ordinals is not None - else len(devices) + visible_ordinals[idx] if visible_ordinals is not None else len(devices) ), "name": name, "memory_total_gb": round(mem_total_mb / 1024, 2), diff --git a/studio/backend/utils/hardware/vram_estimation.py b/studio/backend/utils/hardware/vram_estimation.py index ba1b1dfe61..86069ead3d 100644 --- a/studio/backend/utils/hardware/vram_estimation.py +++ b/studio/backend/utils/hardware/vram_estimation.py @@ -16,9 +16,7 @@ from dataclasses import dataclass, field from typing import Dict, Optional QUANT_4BIT_FACTOR = 16 / 5 -DOUBLE_QUANT_4BIT_FACTOR = ( - 3.6 # bnb_4bit_use_double_quant; see VRAM_ESTIMATION.md section 1 -) +DOUBLE_QUANT_4BIT_FACTOR = 3.6 # bnb_4bit_use_double_quant; see VRAM_ESTIMATION.md section 1 CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3) # calibrated on RTX 5070 Ti NON_FLASH_ATTENTION_FACTOR = ( 12.0 # eager attention score+workspace overhead; see VRAM_ESTIMATION.md section 5 @@ -61,8 +59,7 @@ OPTIMIZER_BYTES_PER_PARAM: Dict[str, int] = { } # (full_ft_multiplier, lora_multiplier) — fraction of num_layers. -# LoRA: frozen base layers skip activation storage, but you always need -# at least ~1 layer in flight during backprop recomputation. +# LoRA: frozen layers skip activation storage, but ~1 is in flight during backprop. GC_LAYER_MULTIPLIERS = { "none": (None, None), "true": (2.0, 1.0), @@ -127,8 +124,7 @@ class VramBreakdown: gradients: int activations: int cuda_overhead: int - # Equals `activations`; retained for backward compatibility with - # consumers that read this field. + # Equals `activations`; kept for backward compat with field consumers. activations_computed: int = 0 @property @@ -143,17 +139,12 @@ class VramBreakdown: ) def min_gpu_vram(self, n_gpus: int) -> int: - """Minimum VRAM a single GPU needs: its shard + non-shardable costs. + """Min VRAM one GPU needs: its shard + non-shardable costs. - Weights/LoRA/optimizer/gradients shard across GPUs. - Activations do NOT shard (the GPU running a layer holds them). + Weights/LoRA/optimizer/gradients shard across GPUs; activations do + NOT (the GPU running a layer holds them). """ - shardable = ( - self.model_weights - + self.lora_adapters - + self.optimizer_states - + self.gradients - ) + shardable = self.model_weights + self.lora_adapters + self.optimizer_states + self.gradients per_gpu_fixed = self.activations + self.cuda_overhead return shardable // max(n_gpus, 1) + per_gpu_fixed @@ -170,7 +161,7 @@ class VramBreakdown: def _first_scalar(value): - # why: ERNIE MoE configs ship moe_intermediate_size / moe_num_experts as + # ERNIE MoE ships moe_intermediate_size / moe_num_experts as # [routed, shared] lists; downstream arithmetic needs the routed scalar. if isinstance(value, (list, tuple)): return value[0] if value else None @@ -178,8 +169,8 @@ def _first_scalar(value): def _max_scalar(value): - # why: Hunyuan-V1-MoE moe_topk can be a per-layer list; activation - # accounting uses the max top-k as a conservative upper bound. + # Hunyuan-V1-MoE moe_topk can be a per-layer list; activation accounting + # uses max top-k as a conservative upper bound. if isinstance(value, (list, tuple)): items = [v for v in value if v is not None] return max(items) if items else None @@ -188,18 +179,15 @@ def _max_scalar(value): def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple: """Layer indices that use dense MLP instead of MoE. Position matters.""" - # why: transformers Exaone-MoE / Laguna / Hy_v3 / GLM-MoE-DSA / GLM4-MoE-Lite / - # Ernie4_5_VL_MoE prefer per-position `mlp_layer_types` over the prefix-style - # `first_k_dense_replace` and may omit `decoder_sparse_step` entirely. + # Exaone-MoE / Laguna / Hy_v3 / GLM-MoE-DSA / GLM4-MoE-Lite / Ernie4_5_VL_MoE + # prefer per-position `mlp_layer_types` over prefix `first_k_dense_replace`. layer_types = getattr(text_config, "mlp_layer_types", None) if layer_types: return tuple( - i - for i, t in enumerate(layer_types[:total_layers]) - if str(t).lower() == "dense" + i for i, t in enumerate(layer_types[:total_layers]) if str(t).lower() == "dense" ) - # why: Llama4TextConfig.__init__ auto-populates self.moe_layers from + # Llama4TextConfig.__init__ auto-populates self.moe_layers from # interleave_moe_layer_step; Llama4TextDecoderLayer dispatches via # `layer_idx in config.moe_layers` (modeling_llama4.py). llama4_moe_layers = getattr(text_config, "moe_layers", None) @@ -207,10 +195,9 @@ def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple: moe_indices = {int(i) for i in llama4_moe_layers} return tuple(i for i in range(total_layers) if i not in moe_indices) - # why: transformers ERNIE 4.5 MoE / ERNIE 4.5 VL MoE declare MoE layers - # via moe_layer_start_index / moe_layer_end_index / moe_layer_interval; - # the model's per-layer guard is `(layer_idx + 1) % interval == 0` with - # start <= layer_idx <= end (modeling_ernie4_5_moe.py). + # ERNIE 4.5 (VL) MoE: layers via moe_layer_start/end_index + interval; + # per-layer guard `(layer_idx+1) % interval == 0` within [start, end] + # (modeling_ernie4_5_moe.py). moe_start = getattr(text_config, "moe_layer_start_index", None) moe_interval = getattr(text_config, "moe_layer_interval", None) if moe_start is not None and moe_interval is not None and int(moe_interval) > 0: @@ -234,9 +221,7 @@ def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple: if sparse_step is not None and sparse_step > 0: mlp_only_set = {int(i) for i in mlp_only} return tuple( - i - for i in range(total_layers) - if i in mlp_only_set or (i + 1) % sparse_step != 0 + i for i in range(total_layers) if i in mlp_only_set or (i + 1) % sparse_step != 0 ) return () @@ -264,8 +249,7 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: intermediate_size = hidden_size * 4 if not all( - v is not None - for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size) + v is not None for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size) ): return None if num_heads <= 0: @@ -273,8 +257,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads) - # why: DBRX places its MoE attrs on the DbrxFFNConfig sub-config; probe - # ffn_config as a secondary source so DBRX is not misclassified as dense. + # DBRX places its MoE attrs on the DbrxFFNConfig sub-config; probe + # ffn_config as a secondary source so DBRX isn't misclassified as dense. ffn_config = getattr(text_config, "ffn_config", None) def _moe_attr(name): @@ -298,8 +282,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: if moe_intermediate_raw is None: moe_intermediate_raw = _moe_attr("ffn_hidden_size") moe_intermediate = _first_scalar(moe_intermediate_raw) - # why: Exaone-MoE / ERNIE families alias num_shared_experts / - # moe_num_shared_experts to the canonical n_shared_experts. + # Exaone-MoE / ERNIE alias num_shared_experts / moe_num_shared_experts + # to the canonical n_shared_experts. n_shared_experts = ( _first_scalar(_moe_attr("n_shared_experts")) or _first_scalar(_moe_attr("num_shared_experts")) @@ -309,9 +293,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: shared_expert_intermediate_size = _moe_attr("shared_expert_intermediate_size") if shared_expert_intermediate_size and n_shared_experts == 0: n_shared_experts = 1 - # why: DBRX exposes moe_top_k, Hunyuan-V1-MoE exposes moe_topk (which can - # be a per-layer list); _max_scalar normalizes list values to the worst - # case so int(...) below cannot crash on the canonical attribute_map path. + # DBRX moe_top_k; Hunyuan-V1-MoE moe_topk (may be a per-layer list). + # _max_scalar normalizes lists to the worst case so int(...) can't crash. num_experts_per_tok = ( _max_scalar(_moe_attr("num_experts_per_tok")) or _max_scalar(_moe_attr("top_k_experts")) @@ -325,14 +308,11 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: dense_layer_indices = _compute_dense_layer_indices(text_config, num_layers) num_dense_layers = len(dense_layer_indices) - # why: Llama4 dense layers use intermediate_size_mlp; routed and shared - # experts use intermediate_size. Llama4TextMoe builds one shared_expert - # per MoE layer (modeling_llama4.py). + # Llama4 dense layers use intermediate_size_mlp; experts use + # intermediate_size. One shared_expert per MoE layer (modeling_llama4.py). intermediate_size_mlp_raw = _first_scalar(_moe_attr("intermediate_size_mlp")) dense_intermediate_size = ( - int(intermediate_size_mlp_raw) - if intermediate_size_mlp_raw is not None - else None + int(intermediate_size_mlp_raw) if intermediate_size_mlp_raw is not None else None ) if ( intermediate_size_mlp_raw is not None @@ -391,9 +371,7 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: None, ) or 0, - quantization_skip_modules = list( - quantization_config.get("llm_int8_skip_modules", []) or [] - ), + quantization_skip_modules = list(quantization_config.get("llm_int8_skip_modules", []) or []), quant_4bit_factor = quant_4bit_factor, moe_has_dense_mlp = bool(getattr(text_config, "enable_moe_block", False)), dense_layer_indices = dense_layer_indices, @@ -402,8 +380,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: def _targets_all_linear(target_modules) -> bool: - # why: peft LoraConfig accepts target_modules="all-linear" as a bare - # string; iterating a string yields chars and never matches the set. + # peft LoraConfig accepts target_modules="all-linear" as a bare string; + # iterating a string yields chars and never matches the set. if isinstance(target_modules, str): target_modules = [target_modules] normalized = {str(module).lower().replace("_", "-") for module in target_modules} @@ -441,10 +419,9 @@ def _is_kv_shared_layer(arch: ModelArchConfig, layer_idx: int) -> bool: if arch.num_kv_shared_layers <= 0: return False first_shared = arch.num_hidden_layers - arch.num_kv_shared_layers - # why: transformers Gemma4 (modeling_gemma4.py:1031, modular_gemma4.py:863) - # uses the same `> 0` guard so a fully-shared config raises during model - # construction; matching upstream avoids producing a detailed estimate - # for a shape the actual model code rejects. + # Gemma4 (modeling_gemma4.py:1031, modular_gemma4.py:863) uses the same + # `> 0` guard so a fully-shared config raises at model construction; + # matching upstream avoids estimating a shape the model code rejects. return layer_idx >= first_shared > 0 @@ -455,9 +432,9 @@ def _is_dense_mlp_layer(arch: ModelArchConfig, layer_idx: int) -> bool: def _per_layer_input_quantizable(arch: ModelArchConfig) -> int: - # why: Gemma4 PLE block adds per_layer_model_projection (single Linear), + # Gemma4 PLE block adds per_layer_model_projection (single Linear), # per_layer_input_gate (per layer), and per_layer_projection (per layer); - # see transformers gemma4/modular_gemma4.py:1077-1083 and :1247-1253. + # see gemma4/modular_gemma4.py:1077-1083 and :1247-1253. pli = arch.hidden_size_per_layer_input if pli <= 0: return 0 @@ -475,23 +452,13 @@ def _per_layer_input_norm_elements(arch: ModelArchConfig) -> int: return hd * n_layers + pli -def _per_layer_input_lora_params( - arch: ModelArchConfig, - r: int, - target_modules, -) -> int: - # why: Unsloth's get_peft_regex (unsloth_zoo/peft_utils.py) requires module - # names to contain a component tag (mlp/attn/...); PLE module names lack - # any tag, so all-linear training does NOT attach LoRA to them. Only count - # PLE LoRA when the user explicitly names PLE modules. +def _per_layer_input_lora_params(arch: ModelArchConfig, r: int, target_modules) -> int: + # get_peft_regex requires a component tag (mlp/attn/...); PLE names lack + # one, so all-linear skips them. Count PLE LoRA only when named explicitly. pli = arch.hidden_size_per_layer_input if pli <= 0: return 0 - targets = ( - {target_modules} - if isinstance(target_modules, str) - else set(target_modules or []) - ) + targets = {target_modules} if isinstance(target_modules, str) else set(target_modules or []) n_layers = arch.num_hidden_layers hd = arch.hidden_size total = 0 @@ -508,11 +475,7 @@ def _layer_attention_dims(arch: ModelArchConfig, layer_idx: int) -> tuple: layer_types = _layer_types(arch) layer_type = layer_types[layer_idx] is_sliding = layer_type == "sliding_attention" - head_dim = ( - arch.global_head_dim - if not is_sliding and arch.global_head_dim - else _head_dim(arch) - ) + head_dim = arch.global_head_dim if not is_sliding and arch.global_head_dim else _head_dim(arch) use_alt_attention = arch.attention_k_eq_v and not is_sliding num_kv_heads = ( arch.num_global_key_value_heads @@ -532,10 +495,7 @@ def _layer_mlp_size(arch: ModelArchConfig, layer_idx: int) -> int: return _dense_mlp_size(arch) -def _text_linear_dims( - arch: ModelArchConfig, - layer_idx: int, -) -> Dict[str, tuple[int, int]]: +def _text_linear_dims(arch: ModelArchConfig, layer_idx: int) -> Dict[str, tuple[int, int]]: hd = arch.hidden_size if _uses_structured_layer_shapes(arch): q_size, kv_size, has_k, has_v = _layer_attention_dims(arch, layer_idx) @@ -574,26 +534,20 @@ def _module_path_matches(skip_module: str, alias: str) -> bool: if alias_parts[0] == "layers": return skip_parts == alias_parts if len(skip_parts) <= len(alias_parts): - # why: transformers BNB quantizer suffix-matches short skip entries - # like ["q_proj"] / ["lm_head"] against full module paths, so a skip - # shorter than the alias is a tail match. + # BNB suffix-matches short skip entries (["q_proj"], ["lm_head"]) so a + # skip shorter than the alias is a tail match. return alias_parts[-len(skip_parts) :] == skip_parts if skip_parts[-len(alias_parts) :] != alias_parts: return False prefix_parts = skip_parts[: len(skip_parts) - len(alias_parts)] if not prefix_parts: return True - # why: bound the prefix to known text-tower roots so VLM skip names like - # vision_tower.model.layers..self_attn.q_proj do not shadow the text - # alias model.layers..self_attn.q_proj. + # Bound the prefix to text-tower roots so VLM skips like + # vision_tower.model.layers... don't shadow the text alias. return ".".join(prefix_parts) in _SKIP_MODULE_TEXT_PREFIXES -def _add_module_aliases( - aliases: Dict[str, str], - canonical: str, - suffix: str, -) -> None: +def _add_module_aliases(aliases: Dict[str, str], canonical: str, suffix: str) -> None: for prefix in ( "", "model", @@ -607,9 +561,7 @@ def _add_module_aliases( aliases[alias] = canonical -def _build_text_module_elements( - arch: ModelArchConfig, -) -> tuple[Dict[str, int], Dict[str, str]]: +def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int], Dict[str, str]]: elements: Dict[str, int] = {} aliases: Dict[str, str] = {} @@ -620,17 +572,12 @@ def _build_text_module_elements( for layer_idx in range(arch.num_hidden_layers): layer_modules: Dict[str, int] = {} dims = _text_linear_dims(arch, layer_idx) - attn_dims = { - name: dim for name, dim in dims.items() if name in ATTENTION_TARGET_MODULES - } - mlp_dims = { - name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES - } + attn_dims = {name: dim for name, dim in dims.items() if name in ATTENTION_TARGET_MODULES} + mlp_dims = {name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES} if is_mla: - # why: _text_linear_dims uses (hd, hd) for q/o; MLA actually splits - # into q_a/q_b/kv_a/kv_b, so emit a single self_attn aggregate at - # the authoritative MLA per-layer total. + # MLA splits q/o into q_a/q_b/kv_a/kv_b; emit a single self_attn + # aggregate at the authoritative MLA per-layer total. layer_modules["self_attn"] = _compute_attn_elements(arch) else: for name, (in_dim, out_dim) in attn_dims.items(): @@ -648,17 +595,13 @@ def _build_text_module_elements( layer_modules["mlp.experts"] = _compute_routed_moe_elements(arch) shared_moe = _compute_shared_moe_elements(arch) if shared_moe: - # why: Qwen3.5-MoE exposes shared expert as - # mlp.shared_expert; Exaone-MoE/Laguna/GLM-style configs use - # mlp.shared_experts. Register both names so child-path - # llm_int8_skip_modules entries match the right shared block. + # Qwen3.5-MoE: mlp.shared_expert; Exaone-MoE/Laguna/GLM: + # mlp.shared_experts. Register both so skip_modules match. layer_modules["mlp.shared_expert"] = shared_moe if arch.moe_has_dense_mlp: - # why: enable_moe_block runs the dense MLP and the MoE - # experts in parallel; register both for skip matching. - # Non-structured _text_linear_dims returns mlp_size from - # _get_mlp_size which prefers moe_intermediate_size, so - # rebuild dense dims from arch.intermediate_size directly. + # enable_moe_block runs dense MLP and experts in parallel; + # register both. Non-structured _get_mlp_size prefers + # moe_intermediate_size, so rebuild dense dims directly. if _uses_structured_layer_shapes(arch): dense_dims = mlp_dims else: @@ -677,15 +620,12 @@ def _build_text_module_elements( ) else: layer_modules.update( - { - f"mlp.{name}": in_dim * out_dim - for name, (in_dim, out_dim) in mlp_dims.items() - } + {f"mlp.{name}": in_dim * out_dim for name, (in_dim, out_dim) in mlp_dims.items()} ) if pli > 0: - # why: register PLE per-layer linears so llm_int8_skip_modules - # entries like model.layers.0.per_layer_input_gate match. + # Register PLE per-layer linears so llm_int8_skip_modules entries + # like model.layers.0.per_layer_input_gate match. layer_modules["per_layer_input_gate"] = hd_global * pli layer_modules["per_layer_projection"] = pli * hd_global @@ -694,20 +634,16 @@ def _build_text_module_elements( for name, value in layer_modules.items() if name == "self_attn" or name.startswith("self_attn.") ) - # why: gemma4 enable_moe_block puts routed experts at the sibling - # layers..experts attribute, not under self.mlp; the layer's "mlp" - # aggregate must reflect only the dense MLP path so a skip module - # `model.layers.0.mlp` does not over-skip into the experts block. + # gemma4 enable_moe_block puts routed experts at sibling + # layers..experts, not under self.mlp; keep the "mlp" aggregate to + # the dense path so a `model.layers.0.mlp` skip doesn't over-skip. is_sibling_experts = bool(arch.moe_has_dense_mlp) mlp_total = sum( value for name, value in layer_modules.items() if ( name == "mlp" - or ( - name.startswith("mlp.") - and not (is_sibling_experts and name == "mlp.experts") - ) + or (name.startswith("mlp.") and not (is_sibling_experts and name == "mlp.experts")) ) ) experts_total = layer_modules.get("mlp.experts", 0) if is_sibling_experts else 0 @@ -730,12 +666,10 @@ def _build_text_module_elements( elements[canonical] = value _add_module_aliases(aliases, canonical, canonical.removeprefix("text.")) if name == "mlp.experts" and arch.moe_has_dense_mlp: - # why: gemma4 enable_moe_block exposes routed experts at - # layers..experts (sibling of self.mlp), not under mlp. + # gemma4: routed experts at sibling layers..experts, not mlp. _add_module_aliases(aliases, canonical, f"layers.{layer_idx}.experts") elif name == "mlp.shared_expert": - # why: Exaone-MoE / Laguna / GLM-style configs use the plural - # `shared_experts` attribute name; register both spellings. + # Exaone-MoE/Laguna/GLM use plural `shared_experts`; add both. _add_module_aliases( aliases, canonical, @@ -764,10 +698,7 @@ def _compute_skipped_quantizable_elements(arch: ModelArchConfig) -> int: pruned = { canonical for canonical in matched - if not any( - canonical != parent and canonical.startswith(f"{parent}.") - for parent in matched - ) + if not any(canonical != parent and canonical.startswith(f"{parent}.") for parent in matched) } return sum(module_elements[canonical] for canonical in pruned) @@ -783,8 +714,8 @@ def _get_mlp_size(arch: ModelArchConfig) -> int: def _dense_mlp_size(arch: ModelArchConfig) -> int: - # why: Llama4 dense layers use intermediate_size_mlp; routed/shared - # experts use intermediate_size. Other configs leave the field None. + # Llama4 dense layers use intermediate_size_mlp; routed/shared experts use + # intermediate_size. Other configs leave the field None. return arch.dense_intermediate_size or arch.intermediate_size @@ -814,7 +745,7 @@ def _compute_dense_mlp_elements(arch: ModelArchConfig) -> int: def _shared_expert_size(arch: ModelArchConfig) -> int: - # why: Qwen3.5-MoE shared expert has its own intermediate_size (default 512) + # Qwen3.5-MoE shared expert has its own intermediate_size (default 512) # distinct from moe_intermediate_size; fall back to routed mlp_size for # families that share it (deepseek-style configs). return arch.shared_expert_intermediate_size or _get_mlp_size(arch) @@ -832,10 +763,8 @@ def _compute_shared_moe_elements(arch: ModelArchConfig) -> int: hd = arch.hidden_size shared_size = _shared_expert_size(arch) total = hd * shared_size * 3 * arch.n_shared_experts - # why: only Qwen2-MoE / Qwen3.5-MoE define a shared_expert_gate Linear - # (hidden_size→1); other families (Exaone-MoE, HY-V3, GLM4-MoE-Lite, Laguna) - # have shared_experts without a gate. shared_expert_intermediate_size is the - # Qwen-style discriminator. + # Only Qwen2/Qwen3.5-MoE add a shared_expert_gate Linear (hidden_size->1); + # shared_expert_intermediate_size is the Qwen-style discriminator. if arch.shared_expert_intermediate_size: total += arch.n_shared_experts * hd return total @@ -874,8 +803,8 @@ def _compute_layer_elements(arch: ModelArchConfig): n_moe = n_layers - n_dense moe_mlp_total = _compute_moe_mlp_elements(arch) * n_moe if arch.moe_has_dense_mlp: - # why: enable_moe_block runs dense MLP and MoE experts in - # parallel; count dense for every layer alongside MoE. + # enable_moe_block runs dense MLP and MoE experts in parallel; + # count dense for every layer alongside MoE. mlp_total = sum(per_layer_dense_mlp) + moe_mlp_total else: dense_only_total = sum( @@ -900,9 +829,7 @@ def _compute_layer_elements(arch: ModelArchConfig): mlp_total = _compute_dense_mlp_elements(arch) * n_layers layernorms = 2 * hd - per_layer_embed = ( - arch.vocab_size_per_layer_input * arch.hidden_size_per_layer_input * n_layers - ) + per_layer_embed = arch.vocab_size_per_layer_input * arch.hidden_size_per_layer_input * n_layers ple_text_linear = _per_layer_input_quantizable(arch) ple_norms = _per_layer_input_norm_elements(arch) embed_tokens = arch.vocab_size * hd + per_layer_embed + ple_norms @@ -911,9 +838,7 @@ def _compute_layer_elements(arch: ModelArchConfig): def compute_model_weights_bytes( - arch: ModelArchConfig, - training_method: str, - load_in_4bit: bool, + arch: ModelArchConfig, training_method: str, load_in_4bit: bool ) -> int: total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch) n_layers = arch.num_hidden_layers @@ -926,9 +851,7 @@ def compute_model_weights_bytes( ) quantized = total_quantizable - skipped_quantizable return int( - quantized * 2 / arch.quant_4bit_factor - + skipped_quantizable * 2 - + non_quantizable * 2 + quantized * 2 / arch.quant_4bit_factor + skipped_quantizable * 2 + non_quantizable * 2 ) return int((total_quantizable + non_quantizable) * 2) @@ -940,11 +863,7 @@ def compute_total_params(arch: ModelArchConfig) -> int: return total_quantizable + layernorms * n_layers + embed_tokens + lm_head -def _lora_attn_elements( - arch: ModelArchConfig, - r: int, - target_modules: list, -) -> int: +def _lora_attn_elements(arch: ModelArchConfig, r: int, target_modules: list) -> int: hd = arch.hidden_size if arch.q_lora_rank is not None: # MLA: q_proj->q_b, k_proj->kv_a, v_proj->kv_b, o_proj->o @@ -974,11 +893,7 @@ def _lora_attn_elements( def _lora_mlp_elements( - hd: int, - mlp_size: int, - r: int, - target_modules: list, - expert_mult: int, + hd: int, mlp_size: int, r: int, target_modules: list, expert_mult: int ) -> int: module_ab = { "gate_proj": (hd * r, r * mlp_size), @@ -992,11 +907,7 @@ def _lora_mlp_elements( return total -def compute_lora_params( - arch: ModelArchConfig, - lora_rank: int, - target_modules: list, -) -> int: +def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: list) -> int: all_linear = _targets_all_linear(target_modules) selected_modules = list(DEFAULT_TARGET_MODULES) if all_linear else target_modules hd = arch.hidden_size @@ -1026,11 +937,10 @@ def compute_lora_params( if n_experts > 1: n_dense = arch.num_dense_layers n_moe = n_layers - n_dense - # why: peft "all-linear" attaches LoRA to nn.Linear only; - # routed experts are nn.Parameter and need explicit - # gate_proj/up_proj/down_proj naming via Unsloth's - # get_moe_target_parameters. Shared experts are nn.Linear and - # are picked up by get_peft_regex. + # peft "all-linear" attaches LoRA to nn.Linear only; routed experts + # are nn.Parameter and need explicit gate_proj/up_proj/down_proj + # naming via Unsloth's get_moe_target_parameters. Shared experts are + # nn.Linear, picked up by get_peft_regex. routed_moe = ( 0 if all_linear @@ -1051,7 +961,7 @@ def compute_lora_params( ) moe_mlp = routed_moe + shared_moe if arch.moe_has_dense_mlp: - # why: parallel dense MLP coexists with MoE on every layer. + # Parallel dense MLP coexists with MoE on every layer. mlp_total = structured_dense_mlp + moe_mlp * n_moe else: dense_only = sum( @@ -1062,16 +972,12 @@ def compute_lora_params( mlp_total = moe_mlp * n_moe + dense_only else: mlp_total = structured_dense_mlp - return ( - attn_total - + mlp_total - + _per_layer_input_lora_params(arch, r, target_modules) - ) + return attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules) elif n_experts > 1: attn_total = _lora_attn_elements(arch, r, selected_modules) * n_layers n_dense = arch.num_dense_layers n_moe = n_layers - n_dense - # why: routed and shared experts may use different intermediate sizes + # Routed and shared experts may use different intermediate sizes # (Qwen3.5-MoE: routed mlp_size != shared_expert_intermediate_size). # See structured branch for the all-linear exclusion rationale; only # routed (nn.Parameter) experts are excluded under all-linear. @@ -1118,9 +1024,7 @@ def compute_lora_params( * n_layers ) - return ( - attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules) - ) + return attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules) def compute_lora_adapter_bytes(lora_params: int) -> int: @@ -1138,16 +1042,13 @@ def compute_gradient_bytes(trainable_params: int) -> int: def _is_linear_attention(attention_implementation: Optional[str]) -> bool: - # why: PyTorch SDPA dispatches to flash/memory-efficient O(n) backends; only + # PyTorch SDPA dispatches to flash/memory-efficient O(n) backends; only # eager (and other non-flash impls) need the quadratic correction. return attention_implementation in LINEAR_ATTENTION_IMPLS def _compute_non_flash_attention_bytes( - arch: ModelArchConfig, - batch_size: int, - seq_len: int, - effective_layers: float, + arch: ModelArchConfig, batch_size: int, seq_len: int, effective_layers: float ) -> int: score_elements = batch_size * arch.num_attention_heads * seq_len * seq_len return int(score_elements * 2 * NON_FLASH_ATTENTION_FACTOR * effective_layers) @@ -1158,17 +1059,16 @@ def _layer_qkv_mlp_sizes(arch: ModelArchConfig, layer_idx: int) -> tuple: is_moe_layer = n_experts > 1 and not _is_dense_mlp_layer(arch, layer_idx) if _uses_structured_layer_shapes(arch): q_size, kv_size, _has_k, _has_v = _layer_attention_dims(arch, layer_idx) - # why: KV-shared layers (Gemma4/Gemma3n) drop k_proj/v_proj WEIGHTS but - # the donor layer's K/V tensors stay alive across the shared range, so - # activation memory still pays for kv_size; only the weight path uses - # has_k/has_v. + # KV-shared layers (Gemma4/Gemma3n) drop k/v WEIGHTS but the donor's + # K/V tensors stay alive, so activations still pay kv_size; only the + # weight path uses has_k/has_v. layer_type = _layer_types(arch)[layer_idx] use_alt_attention = arch.attention_k_eq_v and layer_type != "sliding_attention" kv_count = 1 if use_alt_attention else 2 qkv_size = q_size + kv_size * kv_count if is_moe_layer: - # why: each token routes through `num_experts_per_tok` experts; their - # gate/up/down intermediates are all live during MLP forward. + # Each token routes through num_experts_per_tok experts; all their + # gate/up/down intermediates are live during MLP forward. mlp_size = _get_mlp_size(arch) * arch.num_experts_per_tok if arch.n_shared_experts: mlp_size += _shared_expert_size(arch) * arch.n_shared_experts @@ -1190,23 +1090,17 @@ def _layer_qkv_mlp_sizes(arch: ModelArchConfig, layer_idx: int) -> tuple: def _per_layer_activation_bytes( - arch: ModelArchConfig, - layer_idx: int, - batch_size: int, - seq_len: int, + arch: ModelArchConfig, layer_idx: int, batch_size: int, seq_len: int ) -> int: qkv_size, mlp_size = _layer_qkv_mlp_sizes(arch, layer_idx) activation_qkv = seq_len * batch_size * qkv_size residual_memory = (seq_len * batch_size) * 2 activation_mlp = seq_len * batch_size * (mlp_size + mlp_size) - # why: per_layer_input_gate (hd-sized) and per_layer_projection (pli-sized) - # outputs materialize once per decoder layer when hidden_size_per_layer_input - # is set; see gemma4/modular_gemma4.py:1141-1145. + # PLE gate (hd) + projection (pli) outputs materialize once per decoder + # layer when hidden_size_per_layer_input is set (gemma4 modular:1141-1145). pli = arch.hidden_size_per_layer_input activation_ple = seq_len * batch_size * (arch.hidden_size + pli) if pli > 0 else 0 - return int( - (activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25 - ) + return int((activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25) def compute_activation_bytes( @@ -1227,19 +1121,17 @@ def compute_activation_bytes( if gc_multiplier is None: effective_layers = n_layers linear_bytes = sum( - _per_layer_activation_bytes(arch, i, batch_size, seq_len) - for i in range(n_layers) + _per_layer_activation_bytes(arch, i, batch_size, seq_len) for i in range(n_layers) ) else: effective_layers = gc_multiplier max_layer_bytes = max( - _per_layer_activation_bytes(arch, i, batch_size, seq_len) - for i in range(n_layers) + _per_layer_activation_bytes(arch, i, batch_size, seq_len) for i in range(n_layers) ) linear_bytes = int(max_layer_bytes * effective_layers) - # why: gemma4 per_layer_model_projection runs once outside the per-decoder - # loop and materializes a [B, S, L, PLI] tensor; see modular_gemma4.py:1247. + # gemma4 per_layer_model_projection runs once outside the per-decoder loop + # and materializes a [B, S, L, PLI] tensor; see modular_gemma4.py:1247. pli = arch.hidden_size_per_layer_input if pli > 0: linear_bytes += int(seq_len * batch_size * n_layers * pli * 2 * 1.25) @@ -1257,10 +1149,7 @@ def compute_activation_bytes( ) -def estimate_training_vram( - arch: ModelArchConfig, - config: TrainingVramConfig, -) -> VramBreakdown: +def estimate_training_vram(arch: ModelArchConfig, config: TrainingVramConfig) -> VramBreakdown: method = config.training_method.lower() is_lora = method in ("qlora", "lora") load_in_4bit = config.load_in_4bit or method == "qlora" diff --git a/studio/backend/utils/helper_precache_settings.py b/studio/backend/utils/helper_precache_settings.py new file mode 100644 index 0000000000..db19a2d028 --- /dev/null +++ b/studio/backend/utils/helper_precache_settings.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted opt-in controls for Helper LLM startup pre-cache.""" + +from __future__ import annotations + +import os +from typing import Any + +HELPER_PRECACHE_SETTING_KEY = "helper_model_preload_on_startup" +DEFAULT_HELPER_PRECACHE_ENABLED = False + + +def _coerce_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off", ""}: + return False + return None + + +def helper_model_disabled_by_env() -> bool: + """Return True when existing broad helper-disable env var is active.""" + return os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in {"1", "true"} + + +def get_helper_precache_enabled() -> bool: + """Read the persisted startup pre-cache preference. + + Missing or unreadable settings default to False so Studio startup never + performs optional network work unless the user explicitly opted in. + """ + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(HELPER_PRECACHE_SETTING_KEY, None) + except Exception: + stored = None + parsed = _coerce_bool(stored) + return parsed if parsed is not None else DEFAULT_HELPER_PRECACHE_ENABLED + + +def set_helper_precache_enabled(value: Any) -> bool: + """Persist whether Studio should pre-cache the Helper LLM at startup.""" + parsed = _coerce_bool(value) + if parsed is None: + raise ValueError("Helper LLM startup pre-cache must be true or false.") + + from storage.studio_db import upsert_app_settings + + upsert_app_settings({HELPER_PRECACHE_SETTING_KEY: parsed}) + return parsed + + +def should_preload_helper_on_startup() -> bool: + """Gate the startup pre-cache thread. + + The persisted setting is opt-in and the existing broad disable env var wins. + Explicit AI Assist calls do not use this gate; they remain user-triggered. + """ + return get_helper_precache_enabled() and not helper_model_disabled_by_env() diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py index 9efc281b0b..05eb08067c 100644 --- a/studio/backend/utils/inference/inference_config.py +++ b/studio/backend/utils/inference/inference_config.py @@ -1,13 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Inference configuration loading utilities. - -This module provides functions to load inference parameters (temperature, top_p, top_k, min_p) -from model YAML configuration files, with fallback to default.yaml. -Includes family-based lookup from inference_defaults.json for GGUF models. -""" +"""Load inference params (temperature, top_p, top_k, min_p) from model YAML, family defaults, or default.yaml.""" from pathlib import Path from typing import Dict, Any, Optional @@ -34,10 +28,7 @@ def _load_family_defaults(): return json_path = ( - Path(__file__).parent.parent.parent - / "assets" - / "configs" - / "inference_defaults.json" + Path(__file__).parent.parent.parent / "assets" / "configs" / "inference_defaults.json" ) try: with open(json_path, "r", encoding = "utf-8") as f: @@ -51,29 +42,22 @@ def _load_family_defaults(): def get_family_inference_params(model_id: str) -> Dict[str, Any]: - """ - Look up recommended inference parameters by model family. + """Look up recommended inference params by model family. - Extracts the model family from the identifier (e.g. "unsloth/Qwen3.5-9B-GGUF" -> "qwen3.5") - and returns the matching parameters from inference_defaults.json. - - Args: - model_id: Model identifier (e.g. "unsloth/Qwen3.5-9B-GGUF") - - Returns: - Dict with inference params, or empty dict if no family match. + Extracts the family from the identifier (e.g. "unsloth/Qwen3.5-9B-GGUF" -> + "qwen3.5") and returns matching params from inference_defaults.json, or {}. """ _load_family_defaults() if not _FAMILY_PATTERNS or not _FAMILY_DEFAULTS: return {} - # Normalize: lowercase, strip org prefix + # Normalize: lowercase, strip org prefix. normalized = model_id.lower() if "/" in normalized: normalized = normalized.split("/", 1)[1] - # Match against patterns (ordered longest-match-first in the JSON) + # Match patterns (ordered longest-match-first in the JSON). for pattern in _FAMILY_PATTERNS: if pattern in normalized: params = _FAMILY_DEFAULTS.get(pattern, {}) @@ -90,14 +74,11 @@ def _has_specific_yaml(model_identifier: str) -> bool: script_dir = Path(__file__).parent.parent.parent defaults_dir = script_dir / "assets" / "configs" / "model_defaults" - # Check the mapping if model_identifier.lower() in _REVERSE_MODEL_MAPPING: return True - # For local filesystem paths (e.g. C:\Users\...\model on Windows), - # normalize backslashes so Path().parts splits correctly on POSIX/WSL, - # then try matching the last 1-2 path components against the registry - # (mirrors the logic in load_model_defaults). + # For local paths, normalize backslashes so Path().parts splits correctly, + # then match the last 1-2 components against the registry (mirrors load_model_defaults). _is_local = is_local_path(model_identifier) _normalized = normalize_path(model_identifier) if _is_local else model_identifier @@ -112,9 +93,7 @@ def _has_specific_yaml(model_identifier: str) -> bool: else: _lookup = model_identifier - # Check for exact filename match (basename for local paths to avoid - # passing absolute paths into rglob which raises - # "Non-relative patterns are unsupported" on Windows). + # Exact filename match (basename for local paths; absolute paths break rglob on Windows). model_filename = _lookup.replace("/", "_") + ".yaml" for config_path in defaults_dir.rglob(model_filename): if config_path.is_file(): @@ -124,30 +103,14 @@ def _has_specific_yaml(model_identifier: str) -> bool: def load_inference_config(model_identifier: str) -> Dict[str, Any]: + """Load inference params for a model. + + Priority: model-specific YAML, then family defaults (inference_defaults.json), + then default.yaml. Returns a dict of temperature/top_p/top_k/min_p/etc. """ - Load inference configuration parameters for a model. - - Priority chain: - 1. Model-specific YAML (if it exists and has inference params) - 2. Family-based defaults from inference_defaults.json - 3. default.yaml fallback - - Args: - model_identifier: Model identifier (e.g., "unsloth/llama-3-8b-bnb-4bit") - - Returns: - Dictionary containing inference parameters: - { - "temperature": float, - "top_p": float, - "top_k": int, - "min_p": float - } - """ - # Load model defaults to get inference parameters model_defaults = load_model_defaults(model_identifier) - # Load default.yaml for fallback values + # default.yaml for fallback values. script_dir = Path(__file__).parent.parent.parent defaults_dir = script_dir / "assets" / "configs" / "model_defaults" default_config_path = defaults_dir / "default.yaml" @@ -161,18 +124,18 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]: except Exception as e: logger.warning(f"Failed to load default.yaml: {e}") - # Family-based defaults from inference_defaults.json + # Family-based defaults from inference_defaults.json. family_params = get_family_inference_params(model_identifier) model_inference = model_defaults.get("inference", {}) - # If the model has its own YAML config, those values take priority over family defaults. - # If it only fell back to default.yaml, family defaults take priority. + # Model's own YAML beats family defaults; if it only fell back to + # default.yaml, family defaults win. has_own_yaml = _has_specific_yaml(model_identifier) def _get_param(key, hardcoded_default): if has_own_yaml: - # Model-specific YAML wins, then family fills gaps, then default.yaml + # Model-specific YAML wins, then family fills gaps, then default.yaml. val = model_inference.get(key) if val is not None and isinstance(val, (int, float)): return val @@ -180,7 +143,7 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]: return family_params[key] return default_inference.get(key, hardcoded_default) else: - # No model-specific YAML: family wins, then default.yaml + # No model-specific YAML: family wins, then default.yaml. if key in family_params: return family_params[key] return default_inference.get(key, hardcoded_default) diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py index 2c781f4a7b..f5fd745334 100644 --- a/studio/backend/utils/llama_cpp_freshness.py +++ b/studio/backend/utils/llama_cpp_freshness.py @@ -13,6 +13,7 @@ from __future__ import annotations import json import os +import re import time from datetime import datetime, timezone from pathlib import Path @@ -38,7 +39,6 @@ def _cache_dir() -> Path: """Lazy import so tests can stub storage_roots.""" try: from utils.paths.storage_roots import cache_root - return cache_root() / "llama_cpp_freshness" except Exception: return Path.home() / ".unsloth" / "studio" / "cache" / "llama_cpp_freshness" @@ -46,7 +46,7 @@ def _cache_dir() -> Path: def read_install_marker(binary_path: Optional[str]) -> Optional[dict]: """Walk up from binary_path to find UNSLOTH_PREBUILT_INFO.json. - None means no marker (source build / custom path) or invalid JSON.""" + None = no marker (source build / custom path) or invalid JSON.""" if not binary_path: return None cached = _marker_cache.get(binary_path) @@ -54,10 +54,7 @@ def read_install_marker(binary_path: Optional[str]) -> Optional[dict]: return cached p = Path(binary_path) marker: Optional[dict] = None - # Cover all _find_llama_server_binary layouts: - # /llama-server (1 up) - # /build/bin/llama-server (3 up, Linux/macOS cmake) - # /build/bin/Release/llama-server.exe (4 up, Windows cmake) + # Cover all _find_llama_server_binary layouts (binary is 1-4 dirs deep): for parent in p.parents[:5]: candidate = parent / _INSTALL_MARKER_NAME if candidate.is_file(): @@ -108,11 +105,18 @@ def _save_disk_cache(repo: str, latest_tag: Optional[str]) -> None: def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]: - """GitHub API call. None on any failure (offline, rate-limited, etc).""" + """Newest published release tag for `repo`, by publish time. + + Resolves "latest" the way install_llama_prebuilt.py does (newest + non-draft/non-prerelease by ``published_at``), NOT via GitHub's + ``/releases/latest`` pointer. That pointer sorts by commit date and can lag + behind the build the installer actually installs, so detection and apply + disagreed -- the cause of the downgrade/sticky banner. None on any failure + (offline, rate-limited, etc).""" import urllib.error import urllib.request - url = f"https://api.github.com/repos/{repo}/releases/latest" + url = f"https://api.github.com/repos/{repo}/releases?per_page=30" headers = { "Accept": "application/vnd.github+json", "User-Agent": "unsloth-studio-freshness-check", @@ -132,13 +136,24 @@ def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]: ) as exc: logger.debug("freshness fetch failed", repo = repo, error = str(exc)) return None - tag = data.get("tag_name") - return tag if isinstance(tag, str) and tag else None + if not isinstance(data, list): + return None + published = [ + r + for r in data + if isinstance(r, dict) + and not r.get("draft") + and not r.get("prerelease") + and isinstance(r.get("tag_name"), str) + and r.get("tag_name") + ] + if not published: + return None + newest = max(published, key = lambda r: r.get("published_at") or "") + return newest["tag_name"] -def latest_published_release( - repo: str, *, force_refresh: bool = False -) -> Optional[str]: +def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]: """Latest release tag for `repo`. Memo + disk-cached (24h TTL). None when offline and never previously cached.""" if not repo: @@ -178,19 +193,58 @@ def _parse_installed_at(value: object) -> Optional[datetime]: return dt +def parse_base_build(tag: object) -> Optional[int]: + """Numeric base build from a release tag. Handles both a plain ``bNNNN`` and + a mix-build tag like ``b9596-mix-`` (anchored at the start, so the mix + suffix doesn't defeat it). None for anything not starting with ``bNNNN``.""" + if not isinstance(tag, str): + return None + m = re.match(r"b(\d+)", tag.strip()) + return int(m.group(1)) if m else None + + +def is_behind(installed: Optional[str], latest: Optional[str]) -> bool: + """Whether `installed` is genuinely behind `latest`, comparing the FULL + release identity (so a mix build can legitimately be the latest) with a + base-build guard so a lagging GitHub /releases/latest can never read as an + update or a downgrade. + + - identical tags -> not behind (clears the sticky banner post-update) + - higher base build on `latest` -> behind; lower -> NOT behind (downgrade guard) + - same base build: a different/new mix -> behind, but a bare ``bNNNN`` never + supersedes a mix build (extra PRs) at that base -> not behind + - non-bNNNN tags -> behind (plain inequality, since they already differ) + """ + if not installed or not latest: + return False + installed, latest = installed.strip(), latest.strip() + if installed == latest: + return False + ib, lb = parse_base_build(installed), parse_base_build(latest) + if ib is None or lb is None: + return True + if lb != ib: + return lb > ib + # Same base build, different tags: offer a mix (latest carries a suffix), but + # never offer a bare base over a mix install at the same base. + return latest != f"b{lb}" + + def check_prebuilt_freshness( binary_path: Optional[str], *, threshold_days: int = STALENESS_THRESHOLD_DAYS, now: Optional[datetime] = None, ) -> dict: - """Returns {has_marker, stale, installed_tag, latest_tag, + """Returns {has_marker, stale, behind, installed_tag, latest_tag, installed_at_utc, age_days, published_repo, threshold_days}. - stale = True iff installed != latest AND age >= threshold. - Fails open on missing data (stale stays False).""" + behind = installed genuinely older than latest (see is_behind). + stale = behind AND age >= threshold. + Fails open on missing data (behind/stale stay False).""" out: dict = { "has_marker": False, "stale": False, + "behind": False, "installed_tag": None, "latest_tag": None, "installed_at_utc": None, @@ -202,16 +256,25 @@ def check_prebuilt_freshness( if not marker: return out out["has_marker"] = True + # Display prefers the normalized base ("tag"); comparison below prefers the + # full "release_tag" -- deliberately opposite fallbacks. out["installed_tag"] = marker.get("tag") or marker.get("release_tag") out["installed_at_utc"] = marker.get("installed_at_utc") out["published_repo"] = marker.get("published_repo") + # The marker records both a normalized base tag ("tag", e.g. b9596) and the + # full release tag ("release_tag", e.g. b9596-mix-). Compare against the + # FULL identity, since GitHub /releases/latest returns the full tag_name -- + # comparing the normalized base against the full latest is what produced the + # permanent "downgrade" banner on every mix release. + installed_full = marker.get("release_tag") or marker.get("tag") repo = out["published_repo"] - if not repo or not out["installed_tag"]: + if not repo or not installed_full: return out latest = latest_published_release(repo) out["latest_tag"] = latest - if not latest or latest == out["installed_tag"]: + out["behind"] = is_behind(installed_full, latest) + if not out["behind"]: return out installed_at = _parse_installed_at(out["installed_at_utc"]) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py new file mode 100644 index 0000000000..654ade6cd4 --- /dev/null +++ b/studio/backend/utils/llama_cpp_update.py @@ -0,0 +1,560 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""In-app llama.cpp prebuilt update. + +Builds on utils.llama_cpp_freshness (which detects whether a newer prebuilt +release exists) and adds the *apply* half: run install_llama_prebuilt.py to +download the newest bundle for this host and atomically swap it in place, so +the next model load uses it. + +Design notes: +- Detection is delegated to check_prebuilt_freshness(). We surface an + ``update_available`` flag (installed_tag != latest_tag) which is laxer than + freshness' ``stale`` (which additionally requires the install to be >= 3 days + old). The UI shows the "Update llama.cpp" affordance on update_available. +- The install is slow (download + extract + validate), so it runs on a daemon + thread; callers poll get_update_status() for the job state. +- Everything fails open: a missing marker / offline GitHub / source build just + reports update_available=False and never blocks the app. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Optional + +import structlog + +from utils.llama_cpp_freshness import ( + _INSTALL_MARKER_NAME, + check_prebuilt_freshness, + latest_published_release, + read_install_marker, + reset_caches, +) + +logger = structlog.get_logger(__name__) + +DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" +_INSTALL_TIMEOUT_SECONDS = 1800 # 30 min ceiling for download + build/validate + +# Background job state. Single in-flight update at a time, guarded by _job_lock. +_JOB_IDLE = "idle" +_JOB_RUNNING = "running" +_JOB_SUCCESS = "success" +_JOB_ERROR = "error" + +_job_lock = threading.Lock() +_job: dict = { + "state": _JOB_IDLE, + "message": "", + "from_tag": None, + "to_tag": None, + "error": None, + "progress": None, + "started_at": None, + "finished_at": None, +} + +# Matches the installer's download progress lines, e.g. +# "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s". +_PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(") +# The download dominates the update; extract/validate fill the last slice. +_DOWNLOAD_PROGRESS_CEILING = 0.95 + + +def _utcnow() -> str: + return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + + +def _find_binary() -> Optional[str]: + """Locate the active llama-server binary via the inference backend's own + resolver, so update targets exactly what Studio runs. Lazy import keeps the + heavy inference module off this module's import path.""" + try: + from core.inference.llama_cpp import LlamaCppBackend + return LlamaCppBackend._find_llama_server_binary() + except Exception as exc: # pragma: no cover - defensive + logger.debug("llama update: binary discovery failed", error = str(exc)) + return None + + +def _install_dir_for(binary_path: Optional[str]) -> Optional[Path]: + """The directory holding UNSLOTH_PREBUILT_INFO.json -- i.e. the install root + install_llama_prebuilt.py wrote and the one we re-install into. Walks up from + the binary the same way read_install_marker() does.""" + if not binary_path: + return None + p = Path(binary_path) + for parent in p.parents[:5]: + if (parent / _INSTALL_MARKER_NAME).is_file(): + return parent + return None + + +def _installer_script() -> Optional[Path]: + """Locate install_llama_prebuilt.py. Honours UNSLOTH_LLAMA_INSTALLER, then + searches up from this file for both ``/install_llama_prebuilt.py`` and + ``/studio/install_llama_prebuilt.py`` so it works in the dev tree and + in an installed Studio layout.""" + env = os.environ.get("UNSLOTH_LLAMA_INSTALLER") + if env and Path(env).is_file(): + return Path(env) + here = Path(__file__).resolve() + for up in here.parents: + for cand in (up / "install_llama_prebuilt.py", up / "studio" / "install_llama_prebuilt.py"): + if cand.is_file(): + return cand + return None + + +# Markerless (source-build) installs have no UNSLOTH_PREBUILT_INFO.json, so we +# ask the installer whether an official prebuilt now exists for this host. Memo +# is 24h; only successful answers are cached so a network blip retries. +_RESOLVE_TTL_SECONDS = 24 * 60 * 60 +_resolve_memo: dict = {} + + +def _resolve_prebuilt_for_host(*, force_refresh: bool = False) -> Optional[dict]: + """Run install_llama_prebuilt.py --resolve-prebuilt (no download) and return + {prebuilt_available, repo, release_tag, llama_tag, asset, install_kind} or + None. Fail-open: any error -> None so a source build never blocks the app.""" + now = time.time() + if not force_refresh and _resolve_memo: + if now - _resolve_memo.get("at", 0.0) < _RESOLVE_TTL_SECONDS: + return _resolve_memo.get("value") + script = _installer_script() + if script is None: + return None + value: Optional[dict] = None + try: + proc = subprocess.run( + [ + sys.executable, + str(script), + "--resolve-prebuilt", + "latest", + "--output-format", + "json", + ], + capture_output = True, + text = True, + timeout = 60, + ) + out = (proc.stdout or "").strip() + if proc.returncode == 0 and out: + parsed = json.loads(out.splitlines()[-1]) + if isinstance(parsed, dict): + value = parsed + except Exception as exc: # pragma: no cover - subprocess/json defensive + logger.debug("llama update: resolve-prebuilt failed", error = str(exc)) + value = None + if value is not None: # cache real answers; let failures retry next poll + _resolve_memo.update(at = now, value = value) + return value + + +def _installed_build_number(binary: Optional[str]) -> Optional[int]: + """Best-effort build number from ``llama-server --version`` (e.g. + 'version: 9585 (abc)'). None when unparseable or <= 1: a source build with + no git tags reports 'version: 1', which we treat as unknown (offer update).""" + if not binary: + return None + try: + proc = subprocess.run([binary, "--version"], capture_output = True, text = True, timeout = 20) + except Exception: # pragma: no cover - defensive + return None + m = re.search(r"version:\s*(\d+)", (proc.stderr or "") + (proc.stdout or "")) + if not m: + return None + n = int(m.group(1)) + return n if n > 1 else None + + +def _is_under(path: Path, root: Path) -> bool: + try: + p, r = path.resolve(), root.resolve() + except (OSError, ValueError): + p, r = path, root + return p == r or r in p.parents + + +def _llama_install_root(binary: Optional[str]) -> Optional[Path]: + """The Studio-managed llama.cpp root the active binary lives under, or None + when the binary is unmanaged. Installing anywhere the active binary is not + would not replace what _find_llama_server_binary runs (which prefers a pinned + LLAMA_SERVER_PATH, then UNSLOTH_LLAMA_CPP_PATH, then a llama.cpp tree), so we + refuse rather than silently install into an inactive or foreign tree.""" + marked = _install_dir_for(binary) + if marked is not None: + return marked + if not binary: + return None + # LLAMA_SERVER_PATH is an explicit user pin that always wins in discovery; + # never auto-replace its tree (even a user's own llama.cpp checkout). + if os.environ.get("LLAMA_SERVER_PATH"): + return None + p = Path(binary) + env = os.environ.get("UNSLOTH_LLAMA_CPP_PATH") + if env and _is_under(p, Path(env)): + return Path(env) + for parent in p.parents: + if parent.name == "llama.cpp": + return parent + # PATH / system / custom install: not a managed tree, so do not offer. + return None + + +def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: + """Update status for a markerless (source-build) install: offer the official + prebuilt when one exists for this host and is newer than the installed + binary. None -> caller falls through to the no-marker default (unsupported).""" + res = _resolve_prebuilt_for_host(force_refresh = force_refresh) + if not res or not res.get("prebuilt_available"): + return None + # llama_tag is the upstream build (bNNNN, what --version reports); release_tag + # can be a fork wrapper tag, so compare/display against llama_tag. + latest = res.get("llama_tag") or res.get("release_tag") + if not latest: + return None + # No resolvable install root (e.g. a pinned LLAMA_SERVER_PATH we cannot + # manage) means an apply would not take effect, so do not offer. + if _llama_install_root(binary) is None: + return None + installed_build = _installed_build_number(binary) + m = re.search(r"(\d+)", latest) + latest_build = int(m.group(1)) if m else None + # Suppress only when the source build is reliably newer/equal; unknown + # version (the involuntary source-build case) is treated as behind. + update_available = ( + installed_build is None or latest_build is None or installed_build < latest_build + ) + with _job_lock: + job = dict(_job) + return { + "supported": True, + "update_available": update_available, + "stale": False, + "installed_tag": (f"b{installed_build}" if installed_build else None), + "latest_tag": latest, + "published_repo": res.get("repo"), + "installed_at_utc": None, + "age_days": None, + "source_build": True, + "job": job, + } + + +def get_update_status(*, force_refresh: bool = False) -> dict: + """Report whether a newer prebuilt exists plus the current job state. + + force_refresh bypasses the 24h release cache for an explicit "check now". + """ + binary = _find_binary() + marker = read_install_marker(binary) + + with _job_lock: + job_running = _job["state"] == _JOB_RUNNING + + # No marker = source build / custom path. Offer the official prebuilt if one + # now exists for this host (this is why macOS source builds showed no button). + # Skipped while the updater swaps the tree: each 3s poll would exec the + # half-replaced binary (on Windows that exec can make the installer's + # os.replace fail) and the poller only consumes job progress. + if marker is None and binary is not None and not job_running: + src = _source_build_status(binary, force_refresh = force_refresh) + if src is not None: + return src + + repo = (marker or {}).get("published_repo") or DEFAULT_PUBLISHED_REPO + + if force_refresh and repo: + # Prime the cache so the freshness read below sees the newest tag. + try: + latest_published_release(repo, force_refresh = True) + except Exception as exc: # pragma: no cover - network defensive + logger.debug("llama update: force refresh failed", error = str(exc)) + + freshness = check_prebuilt_freshness(binary) + installed = freshness.get("installed_tag") + latest = freshness.get("latest_tag") + # `behind` compares the full release identity with a base-build guard, so a + # lagging /releases/latest or a mix-tagged latest can't show a false update + # (see llama_cpp_freshness.is_behind). + update_available = bool(freshness.get("has_marker") and freshness.get("behind")) + + with _job_lock: + job = dict(_job) + + return { + "supported": bool(freshness.get("has_marker")), + "update_available": update_available, + "stale": bool(freshness.get("stale")), + "installed_tag": installed, + "latest_tag": latest, + "published_repo": freshness.get("published_repo") or repo, + "installed_at_utc": freshness.get("installed_at_utc"), + "age_days": freshness.get("age_days"), + "source_build": False, + "job": job, + } + + +def _rocm_install_args(asset: Optional[str]) -> list[str]: + """Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh. + The installer probe can miss the gfx arch on amd-smi-only hosts; lemonade + bundles carry the family in the name (rocm-gfx110X), fork bundles only rocm/hip.""" + if not asset: + return [] + low = asset.lower() + if "rocm" not in low and "hip" not in low: + return [] + gfx = re.search(r"-gfx[0-9a-z]+", low) + if gfx: + # _normalize_forwarded_gfx accepts the family form (gfx110x -> gfx110X). + return ["--rocm-gfx", gfx.group(0).lstrip("-")] + return ["--has-rocm"] + + +def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path) -> None: + """Worker: put the backend into a maintenance state, run the installer for + the latest prebuilt, then refresh caches so the next load uses the new build.""" + backend = None + model_was_active = False + try: + # Maintenance state so no load starts a server from the half-swapped binary + # (and the old binary is freed for the swap). Fails open without a backend. + try: + from routes.inference import get_llama_cpp_backend + backend = get_llama_cpp_backend() + except Exception as exc: + logger.debug( + "llama update: backend unavailable, skipping load coordination", error = str(exc) + ) + backend = None + + if backend is not None: + try: + with backend._serial_load_lock: + backend._llama_update_in_progress = True + # is_active covers the loading/unhealthy window is_loaded misses + # (a live process also locks the exe on Windows during the swap). + if getattr(backend, "is_active", False): + model_was_active = True + backend.unload_model() + except Exception as exc: + logger.debug("llama update: load coordination failed", error = str(exc)) + + cmd = [ + sys.executable, + str(script), + "--install-dir", + str(install_dir), + "--llama-tag", + "latest", + "--published-repo", + repo, + ] + cmd.extend(_rocm_install_args(asset)) + logger.info("llama update: installing", cmd = " ".join(cmd)) + # Stream the installer output so download percent lines feed + # job["progress"]; finer milestones via UNSLOTH_PROGRESS_PERCENT_STEP. + env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") + proc = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = env, + ) + timed_out = threading.Event() + + def _kill_on_timeout() -> None: + timed_out.set() + proc.kill() + + watchdog = threading.Timer(_INSTALL_TIMEOUT_SECONDS, _kill_on_timeout) + watchdog.daemon = True + watchdog.start() + tail_lines: list[str] = [] + try: + assert proc.stdout is not None + for line in proc.stdout: + tail_lines.append(line) + if len(tail_lines) > 80: + del tail_lines[0] + m = _PROGRESS_LINE_RE.search(line) + if m is None: + continue + fraction = min(float(m.group(1)) / 100.0, 1.0) * _DOWNLOAD_PROGRESS_CEILING + with _job_lock: + _job["progress"] = max(_job.get("progress") or 0.0, fraction) + returncode = proc.wait() + finally: + watchdog.cancel() + if timed_out.is_set(): + raise RuntimeError(f"installer timed out after {_INSTALL_TIMEOUT_SECONDS}s") + if returncode != 0: + tail = "".join(tail_lines).strip()[-1500:] + raise RuntimeError(f"installer exited {returncode}: {tail or 'no output'}") + + # New UNSLOTH_PREBUILT_INFO.json is on disk; drop in-memory caches and + # re-prime the 24h disk freshness cache with the true newest, so the + # banner can't linger on a stale same-base value after the swap. + reset_caches() + try: + latest_published_release(repo, force_refresh = True) + except Exception as exc: # pragma: no cover - network defensive + logger.debug("llama update: post-install freshness refresh failed", error = str(exc)) + new_marker = read_install_marker(_find_binary()) + new_tag = (new_marker or {}).get("tag") or (new_marker or {}).get("release_tag") + + with _job_lock: + _job.update( + state = _JOB_SUCCESS, + message = ( + f"Updated llama.cpp to {new_tag}." + + (" Reload your model to use it." if model_was_active else "") + ), + to_tag = new_tag, + error = None, + progress = 1.0, + finished_at = _utcnow(), + ) + logger.info("llama update: success", to_tag = new_tag) + except Exception as exc: + logger.warning("llama update: failed", error = str(exc)) + with _job_lock: + _job.update( + state = _JOB_ERROR, + message = "llama.cpp update failed.", + error = str(exc), + finished_at = _utcnow(), + ) + finally: + # Lift the maintenance state so model loads work again, success or not. + if backend is not None: + try: + backend._llama_update_in_progress = False + except Exception: # pragma: no cover - defensive + pass + + +def start_update() -> dict: + """Kick off a background update. Idempotent: a second call while one is + running returns the in-flight job rather than starting another.""" + binary = _find_binary() + marker = read_install_marker(binary) + script = _installer_script() + if script is None: + return { + "started": False, + "reason": "installer_missing", + "message": "install_llama_prebuilt.py could not be located.", + "job": get_update_status()["job"], + } + + # A job already in flight wins over any freshness re-check below (and skips + # its network call). The final lock block re-checks to close the TOCTOU. + with _job_lock: + if _job["state"] == _JOB_RUNNING: + return {"started": False, "reason": "already_running", "job": dict(_job)} + + if marker: + # Mirror the detection guard: a direct POST or a stale banner must not + # start an install when the latest is not actually newer (force a fresh + # check so a stale 24h cache can't wrongly block a real update either). + status = get_update_status(force_refresh = True) + if not status.get("update_available"): + return { + "started": False, + "reason": "up_to_date", + "message": "The installed llama.cpp build is already at the latest prebuilt.", + "job": status["job"], + } + install_dir = _install_dir_for(binary) + repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO + from_tag = marker.get("tag") or marker.get("release_tag") + asset = marker.get("asset") + else: + # Source build / custom path: only proceed when the same detection logic + # would offer the update (prebuilt exists, install is behind, root is + # manageable), so a direct POST cannot downgrade a newer source build. + src = _source_build_status(binary, force_refresh = True) if binary else None + if src is None: + return { + "started": False, + "reason": "no_prebuilt_available", + "message": ( + "No official llama.cpp prebuilt is available for this host, " + "so the source build cannot be swapped automatically." + ), + "job": get_update_status()["job"], + } + if not src.get("update_available"): + return { + "started": False, + "reason": "up_to_date", + "message": "The installed llama.cpp build is already at or newer than the latest prebuilt.", + "job": get_update_status()["job"], + } + res = _resolve_prebuilt_for_host() + install_dir = _llama_install_root(binary) + repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO + from_tag = None + asset = (res or {}).get("asset") + + if install_dir is None: + return { + "started": False, + "reason": "no_install_dir", + "message": "Could not determine the llama.cpp install directory.", + "job": get_update_status()["job"], + } + + with _job_lock: + if _job["state"] == _JOB_RUNNING: + return {"started": False, "reason": "already_running", "job": dict(_job)} + _job.update( + state = _JOB_RUNNING, + message = "Downloading and installing the latest llama.cpp prebuilt...", + from_tag = from_tag, + to_tag = None, + error = None, + progress = 0.0, + started_at = _utcnow(), + finished_at = None, + ) + job_snapshot = dict(_job) + + thread = threading.Thread( + target = _run_update, + args = (install_dir, repo, asset, script), + name = "llama-cpp-update", + daemon = True, + ) + thread.start() + return {"started": True, "reason": None, "job": job_snapshot} + + +def _reset_job_for_tests() -> None: + """Test-only: return the job tracker to idle.""" + with _job_lock: + _job.update( + state = _JOB_IDLE, + message = "", + from_tag = None, + to_tag = None, + error = None, + progress = None, + started_at = None, + finished_at = None, + ) diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 808e2b012e..74d08ac116 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Model and LoRA configuration handling -""" +"""Model and LoRA configuration handling.""" from .model_config import ( ModelConfig, diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index b6b2e11c2e..5a992926ec 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Checkpoint scanning utilities for discovering training runs and their checkpoints. -""" +"""Checkpoint scanning utilities for discovering training runs and checkpoints.""" import json import structlog @@ -16,11 +14,7 @@ logger = get_logger(__name__) def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]: - """ - Read the training loss from a checkpoint's trainer_state.json. - - Returns the loss from the last log_history entry, or None if unavailable. - """ + """Read loss from the last log_history entry of trainer_state.json, or None.""" trainer_state = checkpoint_path / "trainer_state.json" if not trainer_state.exists(): return None @@ -38,14 +32,13 @@ def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]: def scan_checkpoints( outputs_dir: str = str(outputs_root()), ) -> List[Tuple[str, List[Tuple[str, str, Optional[float]]], dict]]: - """ - Scan outputs folder for training runs and their checkpoints. + """Scan outputs folder for training runs and their checkpoints. Returns: - List of tuples: [(model_name, [(display_name, checkpoint_path, loss), ...], metadata), ...] - metadata keys: base_model, peft_type, lora_rank (all optional) - The first entry in each checkpoint list is the main adapter; its loss is - set to the loss of the last (highest-step) intermediate checkpoint. + [(model_name, [(display_name, checkpoint_path, loss), ...], metadata), ...] + metadata keys (optional): base_model, peft_type, lora_rank. + First checkpoint entry is the main adapter; its loss mirrors the last + (highest-step) intermediate checkpoint. """ models = [] outputs_path = resolve_output_dir(outputs_dir) @@ -65,7 +58,7 @@ def scan_checkpoints( if not (config_file.exists() or adapter_config.exists()): continue - # Extract training metadata from adapter_config.json / config.json + # Training metadata from adapter_config.json / config.json metadata: dict = {} try: if adapter_config.exists(): @@ -77,7 +70,7 @@ def scan_checkpoints( cfg = json.loads(config_file.read_text()) metadata["base_model"] = cfg.get("_name_or_path") - # Detect BNB quantization from config.json (present in both cases) + # Detect BNB quantization from config.json if config_file.exists(): if "cfg" not in dir(): cfg = json.loads(config_file.read_text()) @@ -91,27 +84,25 @@ def scan_checkpoints( except Exception: pass - # Fallback: extract base model name from folder name - # e.g. "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct" + # Fallback: extract base model name from the folder name, e.g. + # "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct" if not metadata.get("base_model"): parts = item.name.rsplit("_", 1) if len(parts) == 2 and parts[1].isdigit(): name_part = parts[0] idx = name_part.find("_") if idx > 0: - metadata["base_model"] = ( - name_part[:idx] + "/" + name_part[idx + 1 :] - ) + metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1 :] else: metadata["base_model"] = name_part - # This is a valid training run + # Valid training run. checkpoints = [] - # Placeholder for the main adapter — loss filled from last checkpoint below + # Main adapter placeholder — loss filled from the last checkpoint below. checkpoints.append((item.name, str(item), None)) - # Scan for intermediate checkpoints (checkpoint-N subdirs) + # Scan for intermediate checkpoints (checkpoint-N subdirs). for sub in sorted(item.iterdir()): if not sub.is_dir() or not sub.name.startswith("checkpoint-"): continue @@ -121,7 +112,7 @@ def scan_checkpoints( loss = _read_checkpoint_loss(sub) checkpoints.append((sub.name, str(sub), loss)) - # Assign the last checkpoint's loss to the main adapter entry + # Assign the last checkpoint's loss to the main adapter entry. if len(checkpoints) > 1: last_checkpoint_loss = checkpoints[-1][2] checkpoints[0] = ( @@ -131,9 +122,7 @@ def scan_checkpoints( ) models.append((item.name, checkpoints, metadata)) - logger.debug( - f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)" - ) + logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)") # Sort by modification time (newest first) models.sort(key = lambda x: Path(x[1][0][1]).stat().st_mtime, reverse = True) diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index 5629bac58b..a2912cc843 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -1,10 +1,9 @@ # 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).""" +"""``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 @@ -47,6 +46,10 @@ _METADATA_CACHE: Dict[_CacheKey, Optional[Dict[str, str]]] = {} _CACHE_LOCK = threading.Lock() _CACHE_MAX_ENTRIES = 4096 +# Separate cache for single bool capability keys (e.g. clip.has_audio_encoder), +# keyed by (file cache key, wanted key). None = key absent / file unreadable. +_BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {} + def _cache_key(path: str) -> Optional[_CacheKey]: try: @@ -61,9 +64,9 @@ def _cache_key(path: str) -> Optional[_CacheKey]: 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.""" + """Return ``general.*`` strings from a GGUF header, or ``None`` if the + file is missing, unreadable, or not a GGUF. ``{}`` means valid but + carrying none of the wanted keys.""" key = _cache_key(path) if key is None: return None @@ -152,9 +155,9 @@ _FIXED_VTYPE_SIZES: Dict[int, int] = { 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.""" + """Advance past one GGUF value. ``f.seek(.., 1)`` past EOF is legal on a + regular file, so truncation is caught on the next read; return False only + for unknown types or sanity-bound overflow.""" if vtype == 8: # STRING slen_bytes = f.read(8) if len(slen_bytes) < 8: @@ -193,6 +196,80 @@ def _skip_gguf_value(f, vtype: int) -> bool: return True +def _parse_gguf_bool(path: str, wanted_key: str) -> Optional[bool]: + """Bool value of ``wanted_key`` (GGUF vtype 7), or ``None`` if absent / + unreadable. Mirrors ``_parse_gguf_header`` for a single bool key.""" + try: + with open(path, "rb") as f: + head = f.read(24) + if len(head) < 24: + return None + magic, _version, _tcount, kv_count = struct.unpack(" 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(" Optional[bool]: + """Cached single-bool-key read, keyed by (path, mtime, size, wanted_key).""" + fkey = _cache_key(path) + if fkey is None: + return None + ckey = (fkey, wanted_key) + with _CACHE_LOCK: + if ckey in _BOOL_CACHE: + return _BOOL_CACHE[ckey] + result = _parse_gguf_bool(path, wanted_key) + with _CACHE_LOCK: + while len(_BOOL_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _BOOL_CACHE.pop(next(iter(_BOOL_CACHE))) + except StopIteration: + break + _BOOL_CACHE[ckey] = result + return result + + +def read_mmproj_audio_capability(path: str) -> Optional[bool]: + """``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's + gemma4ua): ``True``/``False`` if present, ``None`` if absent/unreadable. + Flags audio-input models independently of tokenizer token names.""" + return _read_gguf_bool(path, "clip.has_audio_encoder") + + 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: @@ -204,8 +281,7 @@ def is_mmproj_by_metadata(meta: Optional[Dict[str, str]]) -> Optional[bool]: def pairing_score( - weight_meta: Optional[Dict[str, str]], - mmproj_meta: Optional[Dict[str, str]], + 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.""" diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 993995ee57..f61c210cf6 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Model and LoRA configuration handling -""" +"""Model and LoRA configuration handling.""" from dataclasses import dataclass from typing import Optional, Dict, Any @@ -57,37 +55,36 @@ def _env_offline() -> bool: # ── Model size extraction ──────────────────────────────────── import re as _re -_MODEL_SIZE_RE = _re.compile( - r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE -) -# MoE active-parameter pattern: matches "A3B", "A3.5B", etc. -_ACTIVE_SIZE_RE = _re.compile( - r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE -) +_MODEL_SIZE_RE = _re.compile(r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE) +# MoE active-parameter pattern: "A3B", "A3.5B", etc. +_ACTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE) +# Gemma 3n/4 effective-parameter pattern: "E2B", "E4B" -- the runtime +# footprint (MatFormer + per-layer embeddings), which is the size that +# matters for size-gated policies like sub-3B speculative-decoding fallback. +_EFFECTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])e(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE) def extract_model_size_b(model_id: str) -> float | None: """Extract model size in billions from a model identifier. Prefers MoE active-parameter notation (e.g. ``A3B`` in - ``Qwen3.5-35B-A3B``) over the total parameter count. - Handles both ``B`` (billions) and ``M`` (millions) suffixes. + ``Qwen3.5-35B-A3B``), then Gemma effective-parameter notation + (e.g. ``E2B``), over total params. Handles ``B`` (billions) and + ``M`` (millions) suffixes. """ mid = (model_id or "").lower() - active = _ACTIVE_SIZE_RE.search(mid) - if active: - val = float(active.group(1)) - return val / 1000.0 if active.group(2).lower() == "m" else val - size = _MODEL_SIZE_RE.search(mid) - if not size: - return None - val = float(size.group(1)) - return val / 1000.0 if size.group(2).lower() == "m" else val + # First match wins, in priority order: active > effective > total. + for pattern in (_ACTIVE_SIZE_RE, _EFFECTIVE_SIZE_RE, _MODEL_SIZE_RE): + m = pattern.search(mid) + if m: + val = float(m.group(1)) + return val / 1000.0 if m.group(2).lower() == "m" else val + return None -# Model name mapping: maps all equivalent model names to their canonical YAML config file -# Format: "canonical_model_name.yaml": [list of all equivalent model names] -# Based on the model mapper provided - canonical filename is based on the first model name in the mapper +# Maps equivalent model names to their canonical YAML config file. +# Format: "canonical_model_name.yaml": [equivalent model names]. +# Canonical filename derives from the first model name in each list. MODEL_NAME_MAPPING = { # ── Embedding models ── "unsloth_all-MiniLM-L6-v2.yaml": [ @@ -461,7 +458,7 @@ MODEL_NAME_MAPPING = { ], } -# Reverse mapping for quick lookup: model_name -> canonical_filename +# Reverse lookup: model_name -> canonical_filename _REVERSE_MODEL_MAPPING = {} for canonical_file, model_names in MODEL_NAME_MAPPING.items(): for model_name in model_names: @@ -474,19 +471,16 @@ def load_model_config( token: Optional[str] = None, trust_remote_code: bool = True, ): - """ - Load model config with optional authentication control. - """ + """Load model config with optional authentication control.""" from transformers import AutoConfig if token: - # Explicit token provided - use it return AutoConfig.from_pretrained( model_name, trust_remote_code = trust_remote_code, token = token ) if not use_auth: - # Load without any authentication (for public model checks) + # No auth, for public model checks with without_hf_auth(): return AutoConfig.from_pretrained( model_name, @@ -494,7 +488,7 @@ def load_model_config( token = None, ) - # Use default authentication (cached tokens) + # Default auth (cached tokens) return AutoConfig.from_pretrained( model_name, trust_remote_code = trust_remote_code, @@ -511,18 +505,91 @@ _VLM_MODEL_TYPES = { "internvl_chat", "cogvlm2", "minicpmv", + "gemma4", } +# Audio-only models that share the ForConditionalGeneration suffix +# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration). +_AUDIO_ONLY_MODEL_TYPES = {"csm", "whisper"} + # Pre-computed .venv_t5 paths and backend dir for subprocess version switching. -# Vision check uses 5.5.0 (newest, recognizes all architectures). +# Vision check uses the Gemma 4 5.5 sidecar for existing Gemma 4 architectures. from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402 _VENV_T5_DIR = str(_studio_root() / ".venv_t5_550") _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent) -# Inline script executed in a subprocess with transformers 5.x activated. -# Receives model_name and token via argv, prints JSON result to stdout. -_VISION_CHECK_SCRIPT = r""" + +def _is_vlm(config) -> bool: + architectures = getattr(config, "architectures", None) or [] + model_type = getattr(config, "model_type", None) + if model_type in _AUDIO_ONLY_MODEL_TYPES: + return False + return ( + any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) + or hasattr(config, "vision_config") + or hasattr(config, "img_processor") + or hasattr(config, "image_token_index") + or model_type in _VLM_MODEL_TYPES + ) + + +def _raw_config_has_vision_config( + model_name: str, hf_token: Optional[str] = None +) -> Optional[bool]: + try: + if is_local_path(model_name): + config_path = Path(normalize_path(model_name)).expanduser() / "config.json" + else: + from huggingface_hub import hf_hub_download + config_path = Path( + hf_hub_download( + repo_id = model_name, + filename = "config.json", + token = hf_token, + ) + ) + config = json.loads(config_path.read_text()) + architectures = config.get("architectures") or [] + model_type = config.get("model_type") + if model_type in _AUDIO_ONLY_MODEL_TYPES: + return False + return ( + any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) + or "vision_config" in config + or "img_processor" in config + or "image_token_index" in config + or model_type in _VLM_MODEL_TYPES + ) + except Exception as exc: + logger.warning("Could not read config.json for '%s': %s", model_name, exc) + return None + + +# why: inline _is_vlm and constants are prepended so the subprocess stays +# self-contained and does not import the parent backend module graph. +_VISION_CHECK_INLINE_HELPERS = ( + "_VLM_ARCH_SUFFIXES = " + repr(_VLM_ARCH_SUFFIXES) + "\n" + "_VLM_MODEL_TYPES = " + repr(_VLM_MODEL_TYPES) + "\n" + "_AUDIO_ONLY_MODEL_TYPES = " + repr(_AUDIO_ONLY_MODEL_TYPES) + "\n" + "def _is_vlm(config):\n" + " architectures = getattr(config, 'architectures', None) or []\n" + " model_type = getattr(config, 'model_type', None)\n" + " if model_type in _AUDIO_ONLY_MODEL_TYPES:\n" + " return False\n" + " return (\n" + " any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n" + " or hasattr(config, 'vision_config')\n" + " or hasattr(config, 'img_processor')\n" + " or hasattr(config, 'image_token_index')\n" + " or model_type in _VLM_MODEL_TYPES\n" + " )\n" +) + +# Subprocess script run with transformers 5.x active. Takes model_name and +# token via argv, prints JSON result to stdout. +_VISION_CHECK_SCRIPT = ( + r""" import sys, os, json os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -536,32 +603,20 @@ sys.path.insert(0, venv_t5) if backend_dir not in sys.path: sys.path.insert(0, backend_dir) +""" + + _VISION_CHECK_INLINE_HELPERS + + r""" try: from transformers import AutoConfig + kwargs = {"trust_remote_code": True} if token: kwargs["token"] = token config = AutoConfig.from_pretrained(model_name, **kwargs) - is_vlm = False - if hasattr(config, "architectures"): - is_vlm = any( - x.endswith(("ForConditionalGeneration", "ForVisionText2Text")) - for x in config.architectures - ) - if not is_vlm and hasattr(config, "vision_config"): - is_vlm = True - if not is_vlm and hasattr(config, "img_processor"): - is_vlm = True - if not is_vlm and hasattr(config, "image_token_index"): - is_vlm = True - if not is_vlm and hasattr(config, "model_type"): - vlm_types = {"phi3_v","llava","llava_next","llava_onevision", - "internvl_chat","cogvlm2","minicpmv"} - if config.model_type in vlm_types: - is_vlm = True + is_vlm = _is_vlm(config) - model_type = getattr(config, "model_type", "unknown") + model_type = getattr(config, "model_type", None) archs = getattr(config, "architectures", []) print(json.dumps({"is_vision": is_vlm, "model_type": model_type, "architectures": archs})) @@ -569,21 +624,16 @@ except Exception as exc: print(json.dumps({"error": str(exc)})) sys.exit(1) """ +) -def _is_vision_model_subprocess( - model_name: str, hf_token: Optional[str] = None -) -> Optional[bool]: - """Run is_vision_model check in a subprocess with transformers 5.x. +def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]: + """Run is_vision_model in a subprocess with transformers 5.x. - Same pattern as training/inference workers: spawn a clean subprocess - with .venv_t5/ prepended to sys.path so AutoConfig recognizes newer - architectures (glm4_moe_lite, etc.). - - Returns True/False for definitive results, or None for transient failures - (timeouts, subprocess errors) so callers can decide whether to cache - the result. Subprocess failures are treated as transient because they - can be caused by temporary HF/auth/network issues. + Spawns a clean subprocess with .venv_t5/ on sys.path so AutoConfig + recognizes newer architectures. Returns True/False for definitive results, + or None for transient failures (timeouts, subprocess errors), which are not + cached so they can be retried. """ token_arg = hf_token or "" @@ -643,42 +693,35 @@ def _is_vision_model_subprocess( def _token_fingerprint(token: Optional[str]) -> Optional[str]: - """Return a SHA256 digest of the token for use as a cache key. - - Avoids storing the raw bearer token in process memory as a dict key. - """ + """SHA256 digest of the token for use as a cache key (avoids storing the + raw bearer token in process memory).""" if token is None: return None return hashlib.sha256(token.encode("utf-8")).hexdigest() -# Cache vision detection results per session to avoid repeated subprocess spawns. -# Keyed by (normalized_model_name, token_fingerprint) to handle gated models correctly. -# Only definitive results (True/False from successful detection) are cached; -# transient failures (network errors, timeouts) are NOT cached so they can be retried. +# Cache vision detection per session to avoid repeated subprocess spawns. +# Keyed by (normalized_model_name, token_fingerprint) to handle gated models. +# Only definitive results are cached; transient failures (network, timeouts) +# are NOT cached so they can be retried. _vision_detection_cache: Dict[Tuple[str, Optional[str]], bool] = {} _vision_cache_lock = threading.Lock() def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: """ - Detect vision-language models (VLMs) by checking architecture in config. - Works for fine-tuned models since they inherit the base architecture. + Detect vision-language models (VLMs) via architecture in config. Works for + fine-tuned models since they inherit the base architecture. - For models that require transformers 5.x (e.g. GLM-4.7-Flash), the check - runs in a subprocess with .venv_t5/ activated -- same pattern as the - training and inference workers. - - Results are cached per (model_name, token_fingerprint) for the lifetime of - the process to avoid repeated subprocess spawns and HuggingFace API calls. - Transient failures are not cached so they can be retried on the next call. + Models needing transformers 5.x are checked in a .venv_t5/ subprocess. + Results are cached per (model_name, token_fingerprint) for the process + lifetime; transient failures are not cached so they can be retried. Args: model_name: Model identifier (HF repo or local path) - hf_token: Optional HF token for accessing gated/private models + hf_token: Optional HF token for gated/private models """ - # Normalize model name for cache key to avoid duplicate entries for - # different casings of the same HF repo (e.g. "Org/Model" vs "org/model"). + # Normalize model name so different casings of the same repo share a key try: if is_local_path(model_name): resolved_name = normalize_path(model_name) @@ -693,21 +736,17 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: resolved_name = model_name cache_key = (resolved_name, _token_fingerprint(hf_token)) - # Lock-free fast path for cache hits. Uses a sentinel to distinguish - # "key not found" from "value is False" in a single atomic dict.get() call. + # Lock-free fast path for cache hits. Sentinel distinguishes "key not found" + # from "value is False" in a single atomic dict.get() call. _MISS = object() cached = _vision_detection_cache.get(cache_key, _MISS) if cached is not _MISS: return cached - # Compute outside the lock to avoid serializing long-running detection - # (subprocess spawns with 60s timeout, HF API calls) across all models. - # The tradeoff: two concurrent calls for the same uncached model may - # both run detection, but they produce the same result and the second - # write is a benign no-op. + # Compute outside the lock so long-running detection isn't serialized across + # models. Two concurrent calls may both run, but produce the same result. result = _is_vision_model_uncached(resolved_name, hf_token) - # Only cache definitive results; None means a transient failure occurred - # and we should retry on the next call instead of locking in a wrong answer. + # Only cache definitive results; None is a transient failure, retry later. if result is not None: with _vision_cache_lock: _vision_detection_cache[cache_key] = result @@ -715,20 +754,14 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: return False -def _is_vision_model_uncached( - model_name: str, hf_token: Optional[str] = None -) -> Optional[bool]: - """Uncached vision model detection -- called by is_vision_model(). +def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]: + """Uncached vision detection; use is_vision_model() instead. - Returns True/False for definitive results, or None when detection failed - due to a transient error (network, timeout, subprocess failure) so the - caller knows not to cache the result. - - Do not call directly; use is_vision_model() instead. + Returns True/False for definitive results, or None on transient errors + (network, timeout, subprocess failure) so the caller knows not to cache. """ - # Models that need transformers 5.x must be checked in a subprocess - # because AutoConfig in the main process (transformers 4.57.x) doesn't - # recognize their architectures. + # Models needing transformers 5.x must be checked in a subprocess: the main + # process (transformers 4.57.x) doesn't recognize their architectures. from utils.transformers_version import needs_transformers_5 if needs_transformers_5(model_name): @@ -736,56 +769,36 @@ def _is_vision_model_uncached( "Model '%s' needs transformers 5.x -- checking vision via subprocess", model_name, ) - return _is_vision_model_subprocess(model_name, hf_token = hf_token) + result = _is_vision_model_subprocess(model_name, hf_token = hf_token) + if result is not None: + return result + return _raw_config_has_vision_config(model_name, hf_token = hf_token) try: config = load_model_config(model_name, use_auth = True, token = hf_token) - # Exclude audio-only models that share ForConditionalGeneration suffix + # Exclude audio-only models sharing the ForConditionalGeneration suffix # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration) - _audio_only_model_types = {"csm", "whisper"} model_type = getattr(config, "model_type", None) - if model_type in _audio_only_model_types: + if model_type in _AUDIO_ONLY_MODEL_TYPES: return False - # Check 1: Architecture class name patterns - if hasattr(config, "architectures"): - is_vlm = any(x.endswith(_VLM_ARCH_SUFFIXES) for x in config.architectures) - if is_vlm: - logger.info( - f"Model {model_name} detected as VLM: architecture {config.architectures}" - ) - return True - - # Check 2: Has vision_config (most VLMs: LLaVA, Gemma-3, Qwen2-VL, etc.) - if hasattr(config, "vision_config"): - logger.info(f"Model {model_name} detected as VLM: has vision_config") + if _is_vlm(config): + archs = getattr(config, "architectures", None) or [] + logger.info( + "Model %s detected as VLM (model_type=%s, architectures=%s)", + model_name, + model_type, + archs, + ) return True - # Check 3: Has img_processor (Phi-3.5 Vision uses this instead of vision_config) - if hasattr(config, "img_processor"): - logger.info(f"Model {model_name} detected as VLM: has img_processor") - return True - - # Check 4: Has image_token_index (common in VLMs for image placeholder tokens) - if hasattr(config, "image_token_index"): - logger.info(f"Model {model_name} detected as VLM: has image_token_index") - return True - - # Check 5: Known VLM model_type values that may not match above checks - if hasattr(config, "model_type"): - if config.model_type in _VLM_MODEL_TYPES: - logger.info( - f"Model {model_name} detected as VLM: model_type={config.model_type}" - ) - return True - return False except Exception as e: logger.warning(f"Could not determine if {model_name} is vision model: {e}") - # Permanent failures (model not found, gated, bad config) should be - # cached as False. Transient failures (network, timeout) should not. + # Permanent failures (not found, gated, bad config) cache as False; + # transient ones (network, timeout) should not. try: from huggingface_hub.errors import RepositoryNotFoundError, GatedRepoError except ImportError: @@ -807,14 +820,15 @@ def _is_vision_model_uncached( VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm") -# Cache detection results per session to avoid repeated API calls +# Cache detection per session to avoid repeated API calls _audio_detection_cache: Dict[str, Optional[str]] = {} -# Tokenizer token patterns → audio_type (all 6 types detected from tokenizer_config.json) +# Tokenizer token patterns → audio_type (all 6 types from tokenizer_config.json) _AUDIO_TOKEN_PATTERNS = { "csm": lambda tokens: "<|AUDIO|>" in tokens and "<|audio_eos|>" in tokens, "whisper": lambda tokens: "<|startoftranscript|>" in tokens, - "audio_vlm": lambda tokens: "" in tokens, + # Gemma 3n: ; Gemma 4: <|audio|> (not csm's <|AUDIO|>). + "audio_vlm": lambda tokens: "" in tokens or "<|audio|>" in tokens, "bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens), "dac": lambda tokens: ( "<|audio_start|>" in tokens @@ -822,20 +836,16 @@ _AUDIO_TOKEN_PATTERNS = { and "<|text_start|>" in tokens and "<|text_end|>" in tokens ), - "snac": lambda tokens: ( - sum(1 for t in tokens if t.startswith(" 10000 - ), + "snac": lambda tokens: (sum(1 for t in tokens if t.startswith(" 10000), } def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Optional[str]: - """ - Dynamically detect if a model is an audio model and return its type. + """Detect if a model is an audio model and return its type. - Fully dynamic — works for any model, not just known ones. - Uses tokenizer_config.json special tokens to detect all 6 audio types. - - Returns: audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') or None. + Works for any model via tokenizer_config.json special tokens. + Returns an audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper', + 'audio_vlm') or None. """ if model_name in _audio_detection_cache: return _audio_detection_cache[model_name] @@ -848,13 +858,11 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option return result -def _detect_audio_from_tokenizer( - model_name: str, hf_token: Optional[str] = None -) -> Optional[str]: - """Detect audio type from tokenizer special tokens (for LLM-based audio models). +def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]: + """Detect audio type from tokenizer special tokens. - First checks local HF cache, then fetches tokenizer_config.json from HuggingFace. - Checks added_tokens_decoder for distinctive patterns. + Checks local HF cache first, then fetches tokenizer_config.json from HF; + examines added_tokens_decoder for distinctive patterns. """ def _check_token_patterns(tok_config: dict) -> Optional[str]: @@ -867,7 +875,7 @@ def _detect_audio_from_tokenizer( return audio_type return None - # 1) Check local HF cache first (works for gated/offline models) + # 1) Local HF cache first (works for gated/offline models) try: repo_dir = get_cache_path(model_name) if repo_dir is not None and repo_dir.exists(): @@ -893,7 +901,6 @@ def _detect_audio_from_tokenizer( import os paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"] - # Use provided token, or fall back to env token = hf_token or os.environ.get("HF_TOKEN") headers = {} if token: @@ -912,17 +919,12 @@ def _detect_audio_from_tokenizer( return None except Exception as e: - logger.debug( - f"Could not detect audio type from tokenizer for {model_name}: {e}" - ) + logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}") return None def is_audio_input_type(audio_type: Optional[str]) -> bool: - """Check if an audio_type accepts audio input (ASR/speech understanding). - - Whisper (ASR) and audio_vlm (Gemma3n) accept audio input. - """ + """True if an audio_type accepts audio input: whisper (ASR), audio_vlm (Gemma3n).""" return audio_type in ("whisper", "audio_vlm") @@ -931,8 +933,25 @@ 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``). +def _is_mtp_drafter(path: str) -> bool: + """True for a separate-file MTP drafter (speculative head), a companion + to the main model rather than a selectable quant: the repo-root + ``mtp-*.gguf`` or the ``MTP/`` subdir copies (Gemma 4). + + Mirrors hub.utils.gguf.is_mtp_drafter_path (utils cannot import hub). + Must be excluded everywhere mmproj is, or the drafter leaks into variant + menus (a phantom quant) and quant-matched file lookups -- e.g. a ``Q8_0`` + request must not resolve to ``MTP/...-Q8_0-MTP.gguf``, which sorts ahead + of the real weight. + """ + p = path.lower() + if not p.endswith(".gguf"): + return False + name = p.rsplit("/", 1)[-1] + return name.startswith("mtp-") or "/mtp/" in f"/{p}" + + +# Family tokens for #5347's filename fallback. Lowercase; order irrelevant. _MODEL_FAMILY_TOKENS: tuple[str, ...] = ( "qwen", "gemma", @@ -965,8 +984,8 @@ _MODEL_FAMILY_TOKENS: tuple[str, ...] = ( ) -# Word-bounded match: any letter on either side disqualifies. Stops -# ``phi`` matching ``sapphire``, ``yi`` matching ``tiny``, etc. +# Word-bounded match: a letter on either side disqualifies (stops ``phi`` +# matching ``sapphire``, ``yi`` matching ``tiny``). _FAMILY_TOKEN_RE_CACHE: Dict[str, "_re.Pattern[str]"] = {} @@ -993,8 +1012,8 @@ def _detect_family_token(filename: str) -> Optional[str]: 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.""" + """Launcher guard: 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: @@ -1087,7 +1106,7 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional continue if resolved in seen_resolved: continue - # Prefer ``general.type=='mmproj'``; fall back to filename. + # Prefer ``general.type=='mmproj'``, else 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)): @@ -1139,35 +1158,84 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional return str(best[1]) -def detect_gguf_model(path: str) -> Optional[str]: +def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[str]: + """Find the separate MTP drafter (``mtp-*.gguf``) for a local GGUF model. + + The drafter that pairs with the main weights sits at the repo/snapshot + root (Gemma 4); the weight itself may be at the root or in a quant subdir, + so scan the weight's directory and ``search_root``. Matches by the + ``mtp-`` filename prefix unsloth uses for ``-hf`` auto-discovery -- the + same signal as the HF download path. Repos that bake the head into the + main GGUF (Qwen) have no such sibling, so this returns None. + + Pairs by name so a multi-model folder can't attach a foreign drafter: + unsloth names the drafter ``mtp-.gguf`` where ```` prefixes + the weight filename across all Gemma 4 repos (e.g. + ``mtp-gemma-4-12B-it.gguf`` next to ``gemma-4-12B-it-qat-Q4_0.gguf``). + An unmatched drafter is skipped (fail-safe: no MTP). """ - Check if the given local path is or contains a GGUF model file. + p = Path(path) + weight_name = p.name.lower() if p.suffix.lower() == ".gguf" else None + start_dir = p.parent if p.is_file() else p + dirs = [start_dir] + if search_root is not None: + dirs.append(Path(search_root)) + for d in dirs: + try: + entries = sorted(d.iterdir()) + except OSError: + continue + for f in entries: + name = f.name.lower() + if not (name.startswith("mtp-") and name.endswith(".gguf")): + continue + stem = name[len("mtp-") : -len(".gguf")] + if not stem or (weight_name is not None and not weight_name.startswith(stem)): + continue + try: + if f.is_file(): + return str(f.resolve()) + except OSError: + continue + return None - Handles two cases: - 1. path is a direct .gguf file path - 2. path is a directory containing .gguf files - Skips mmproj (vision projection) files — those must be passed via - ``--mmproj``, not ``-m``. Use :func:`detect_mmproj_file` instead. +def detect_gguf_model(path: str) -> Optional[str]: + """Check if a local path is or contains a GGUF model file. - Returns the full path to the .gguf file if found, None otherwise. - For HuggingFace repo detection, use detect_gguf_model_remote() instead. + Handles a direct .gguf path or a directory of .gguf files. Skips mmproj + files (pass those via ``--mmproj``; see :func:`detect_mmproj_file`). Returns + the .gguf path or None. For HF repos, use detect_gguf_model_remote(). """ p = Path(path) # Case 1: direct .gguf file - if p.suffix.lower() == ".gguf" and p.is_file(): - if _is_mmproj(p.name): + if p.suffix.lower() == ".gguf": + # Companions are not models: rejecting a drafter here also keeps + # detect_mtp_file from pairing the same file with itself + # (-m drafter --model-draft drafter). Include the immediate parent + # dir so the MTP/ subdir copies are caught -- the basename alone + # (...-MTP.gguf) doesn't match the predicate's mtp- prefix. + if _is_mmproj(p.name) or _is_mtp_drafter(f"{p.parent.name}/{p.name}"): return None - # Use absolute (not resolve) to preserve symlink names -- e.g. - # Ollama .studio_links/model.gguf -> blobs/sha256-... should - # keep the readable symlink name, not the opaque blob hash. - return str(p.absolute()) + # Extension is authoritative: don't gate on is_file()/exists(), which + # can fail in the Windows lock window after llama-server is killed. + try: + is_dir = p.is_dir() + except OSError: + is_dir = False # stat() unavailable in the lock window + if not is_dir: + return str(p.absolute()) # absolute() keeps symlink names readable + # Directory named "*.gguf": fall through to the dir scan below. - # Case 2: directory containing .gguf files (skip mmproj) + # Case 2: directory containing .gguf files (skip mmproj / MTP drafter) if p.is_dir(): gguf_files = sorted( - (f for f in _iter_gguf_files(p) if not _is_mmproj(f.name)), + ( + f + for f in _iter_gguf_files(p) + if not _is_mmproj(f.name) and not _is_mtp_drafter(f"{f.parent.name}/{f.name}") + ), key = lambda f: f.stat().st_size, reverse = True, ) @@ -1177,12 +1245,9 @@ def detect_gguf_model(path: str) -> Optional[str]: return None -# Preferred GGUF quantization levels, in descending priority. -# Q4_K_M is a good default: small, fast, acceptable quality. -# UD (Unsloth Dynamic) variants are always preferred over standard quants -# because they provide better quality per bit. If the repo has no UD variants -# (e.g., bartowski repos), the standard quants are used as fallback. -# Ordered by best size/quality tradeoff, not raw quality. +# Preferred GGUF quant levels, descending priority. UD (Unsloth Dynamic) +# variants beat standard quants on quality per bit; repos without UD fall back +# to standard quants. Ordered by size/quality tradeoff, not raw quality. _GGUF_QUANT_PREFERENCE = [ # UD variants (best quality per bit) -- Q4 is the sweet spot "UD-Q4_K_XL", @@ -1226,23 +1291,16 @@ _GGUF_QUANT_PREFERENCE = [ def _pick_best_gguf(filenames: list[str]) -> Optional[str]: - """ - Pick the best GGUF file from a list of filenames. - - Prefers quantization levels in _GGUF_QUANT_PREFERENCE order. - Falls back to the first .gguf file found. - """ + """Pick the best GGUF file: quant levels in _GGUF_QUANT_PREFERENCE order, else first .gguf.""" gguf_files = [f for f in filenames if f.lower().endswith(".gguf")] if not gguf_files: return None - # Try preferred quantization levels for quant in _GGUF_QUANT_PREFERENCE: for f in gguf_files: if quant in f: return f - # Fallback: first GGUF file return gguf_files[0] @@ -1257,7 +1315,7 @@ class GgufVariantInfo: def _extract_quant_label(filename: str) -> str: """ - Extract quantization label like Q4_K_M, IQ4_XS, BF16 from a GGUF filename. + Extract quant label like Q4_K_M, IQ4_XS, BF16 from a GGUF filename. Examples: "gemma-3-4b-it-Q4_K_M.gguf" → "Q4_K_M" @@ -1284,8 +1342,8 @@ def _extract_quant_label(filename: str) -> str: ) match = re.search(quant_re, stem, re.IGNORECASE) # Subdir layouts like ``BF16/foo.gguf`` keep the quant in the directory, - # not the basename. Look at the parent dirs too so the variant label - # matches the snapshot-relative path produced elsewhere. + # not the basename. Check parent dirs too so the label matches the + # snapshot-relative path produced elsewhere. if not match and "/" in filename: parents = filename.rsplit("/", 1)[0] for segment in reversed(parents.split("/")): @@ -1296,16 +1354,16 @@ def _extract_quant_label(filename: str) -> str: if match: prefix = match.group(1) or "" return f"{prefix}{match.group(2)}" - # Fallback: last segment after hyphen + # Fallback: last hyphen-separated segment return stem.split("-")[-1] def _iter_hf_cache_snapshots(repo_id: str): """Yield HF cache snapshot dirs for *repo_id*, newest first. - Empty generator if HF_HUB_CACHE is missing, the repo isn't cached, - or has no snapshots. Repo name match is case-insensitive to handle - casing drift between download time and lookup. + Empty if HF_HUB_CACHE is missing, the repo isn't cached, or has no + snapshots. Repo name match is case-insensitive to handle casing drift + between download time and lookup. """ try: from huggingface_hub import constants as hf_constants @@ -1340,9 +1398,7 @@ def _iter_hf_cache_snapshots(repo_id: str): yield from snap_dirs -def _list_gguf_variants_from_hf_cache( - repo_id: str, -) -> Optional[tuple[list[GgufVariantInfo], bool]]: +def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: """Variants from the local HF cache snapshot, or None if not cached.""" for snap in _iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snap)) @@ -1352,21 +1408,19 @@ def _list_gguf_variants_from_hf_cache( def list_gguf_variants( - repo_id: str, - hf_token: Optional[str] = None, + repo_id: str, hf_token: Optional[str] = None ) -> tuple[list[GgufVariantInfo], bool]: - """ - List all GGUF quantization variants in a HuggingFace repo. + """List all GGUF quant variants in a HF repo. - Separates main model files from mmproj (vision projection) files. - The presence of mmproj files indicates a vision-capable model. + Separates main model files from mmproj (vision projection) files; mmproj + presence flags a vision-capable model. Returns: - (variants, has_vision): list of non-mmproj GGUF variants + vision flag. + (variants, has_vision): non-mmproj GGUF variants + vision flag. """ from huggingface_hub import model_info as hf_model_info - # Offline: skip the API and serve from cache. + # Offline: skip the API and serve from cache if _env_offline(): cached = _list_gguf_variants_from_hf_cache(repo_id) if cached is not None: @@ -1375,9 +1429,9 @@ def list_gguf_variants( try: info = hf_model_info(repo_id, token = hf_token, files_metadata = True) except Exception as e: - # Permanent errors (deleted/gated/bad revision) must surface to - # the caller; serving stale cache here would mask the real cause. - # Matches the early-return in ``detect_gguf_model_remote``. + # Permanent errors (deleted/gated/bad revision) must surface to the + # caller; serving stale cache would mask the real cause. Matches the + # early-return in ``detect_gguf_model_remote``. if type(e).__name__ in ( "RepositoryNotFoundError", "GatedRepoError", @@ -1399,7 +1453,7 @@ def list_gguf_variants( has_vision = False quant_totals: dict[str, int] = {} # quant -> total bytes - quant_first_file: dict[str, str] = {} # quant -> first filename (for display) + quant_first_file: dict[str, str] = {} # quant -> first filename (display) for sibling in info.siblings: fname = sibling.rfilename @@ -1407,10 +1461,13 @@ def list_gguf_variants( continue size = sibling.size or 0 - # mmproj files are vision projection models, not main model files + # mmproj files are vision projections, not main model files if "mmproj" in fname.lower(): has_vision = True continue + # MTP drafters are speculative-decoding companions, not quants. + if _is_mtp_drafter(fname): + continue quant = _extract_quant_label(fname) quant_totals[quant] = quant_totals.get(quant, 0) + size @@ -1426,9 +1483,8 @@ def list_gguf_variants( ) ) - # Sort by size descending (largest = best quality first). - # Recommended pinning and OOM demotion are handled client-side - # where GPU VRAM info is available. + # Sort by size descending (largest = best quality first); pinning and OOM + # demotion happen client-side where GPU VRAM info exists. variants.sort(key = lambda v: -v.size_bytes) return variants, has_vision @@ -1437,11 +1493,10 @@ def list_gguf_variants( def _resolve_gguf_dir(p: Path) -> Optional[Path]: """Resolve a path to the directory containing GGUF variants. - If *p* is already a directory, returns it directly. If *p* is a ``.gguf`` - file whose parent directory has model metadata (``config.json`` or - ``adapter_config.json``), returns the parent -- all GGUFs in that - directory belong to the same model. Returns ``None`` for loose standalone - GGUFs (no config) to avoid cross-wiring unrelated models. + Directory *p* returns directly. A ``.gguf`` file whose parent dir has + model metadata (``config.json`` or ``adapter_config.json``) returns the + parent -- all GGUFs there belong to the same model. Returns ``None`` for + loose standalone GGUFs (no config) to avoid cross-wiring unrelated models. """ if p.is_dir(): return p @@ -1456,17 +1511,14 @@ def _resolve_gguf_dir(p: Path) -> Optional[Path]: return None -def list_local_gguf_variants( - directory: str, -) -> tuple[list[GgufVariantInfo], bool]: - """List GGUF quantization variants in a local directory. +def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], bool]: + """List GGUF quant variants in a local directory. - Mirrors :func:`list_gguf_variants` but reads from the filesystem - instead of the HuggingFace API. Aggregates shard sizes by quant - label so that split GGUFs appear as a single variant. + Like :func:`list_gguf_variants` but reads the filesystem. Aggregates shard + sizes by quant label so split GGUFs appear as one variant. Returns: - (variants, has_vision): list of non-mmproj GGUF variants + vision flag. + (variants, has_vision): non-mmproj GGUF variants + vision flag. """ p = _resolve_gguf_dir(Path(directory)) if p is None: @@ -1476,10 +1528,10 @@ def list_local_gguf_variants( quant_first_file: dict[str, str] = {} has_vision = False - # Recurse so variant-specific subdirectories (e.g. ``BF16/...gguf`` - # used by some HF GGUF repos for the largest quants) are picked up. - # Filenames in the result preserve the relative subpath so that - # ``_find_local_gguf_by_variant`` can locate the file again. + # Recurse so variant-specific subdirs (e.g. ``BF16/...gguf`` used by + # some HF GGUF repos for the largest quants) are picked up. Result + # filenames keep the relative subpath so ``_find_local_gguf_by_variant`` + # can locate the file again. for f in sorted(_iter_gguf_files(p, recursive = True)): if _is_mmproj(f.name): has_vision = True @@ -1488,9 +1540,11 @@ def list_local_gguf_variants( size = f.stat().st_size except OSError: size = 0 - # Pass the relative path so ``BF16/foo.gguf`` and ``Q4_K_M/foo.gguf`` - # produce distinct quant labels instead of collapsing on basename. + # Use the relative path so ``BF16/foo.gguf`` and ``Q4_K_M/foo.gguf`` + # get distinct quant labels instead of collapsing on basename. rel = f.relative_to(p).as_posix() + if _is_mtp_drafter(rel): + continue quant = _extract_quant_label(rel) quant_totals[quant] = quant_totals.get(quant, 0) + size if quant not in quant_first_file: @@ -1511,8 +1565,8 @@ def list_local_gguf_variants( def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: """Find the GGUF file in *directory* matching a quantization *variant*. - For sharded GGUFs (multiple files with the same quant label), returns - the first shard (sorted by name) which is what ``llama-server -m`` expects. + For sharded GGUFs (multiple files sharing a quant label), returns the + first shard (sorted by name), which is what ``llama-server -m`` expects. Returns the resolved absolute path, or ``None`` if no match. """ @@ -1520,14 +1574,15 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: if p is None: return None - # Recurse into subdirectories so variants stored under a quant-named - # subdir (e.g. ``BF16/foo-BF16-00001-of-00002.gguf``) are found. - # Match against the relative path so the quant label can come from - # the directory name when the basename omits it. + # Recurse so variants under a quant-named subdir (e.g. + # ``BF16/foo-BF16-00001-of-00002.gguf``) are found. Match the relative + # path so the quant label can come from the dir name when the basename + # omits it. matches = sorted( f for f in _iter_gguf_files(p, recursive = True) if not _is_mmproj(f.name) + and not _is_mtp_drafter(f.relative_to(p).as_posix()) and _extract_quant_label(f.relative_to(p).as_posix()) == variant ) if matches: @@ -1538,39 +1593,26 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]: """Best GGUF filename for *repo_id* from the local HF cache, or None. - Excludes mmproj (vision projector) files so a partial cache that - only has the projector cannot route the projector as the main model. + Excludes mmproj (vision projector) files so a partial cache holding only + the projector cannot route it as the main model. """ for snap in _iter_hf_cache_snapshots(repo_id): rel_files = [ - f.relative_to(snap).as_posix() + rel for f in _iter_gguf_files(snap, recursive = True) - if not _is_mmproj(f.name) + if not _is_mtp_drafter(rel := f.relative_to(snap).as_posix()) and not _is_mmproj(f.name) ] if rel_files: return _pick_best_gguf(rel_files) return None -def detect_gguf_model_remote( - repo_id: str, - hf_token: Optional[str] = None, -) -> Optional[str]: - """ - Check if a HuggingFace repo contains GGUF files. +def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Optional[str]: + """Return the best GGUF filename in a HF repo, or None. - Returns the filename of the best GGUF file in the repo, or None. - - Retries on transient HF Hub failures (network hiccups, 5xx, slow - cold-start of the API). Without retry, a single transient failure - here returns None silently and the caller treats the repo as - non-GGUF -- which on Apple Silicon (Mac UI route) means falling - through to the MLX backend, which then fails opening a non-existent - config.json on the GGUF-only repo. Three attempts with 1s/2s/4s - backoff covers the typical free-runner HF Hub flakiness. - - When offline, falls back to the local HF cache so a downloaded - repo is still routed to llama-server (not MLX/Unsloth). + Retries (3 attempts, 1s/2s/4s backoff) on transient HF Hub failures: a + silent None would make the caller treat a GGUF-only repo as non-GGUF and + fall through to MLX on Apple Silicon. Offline falls back to the local cache. """ import time from huggingface_hub import model_info as hf_model_info @@ -1588,7 +1630,7 @@ def detect_gguf_model_remote( return _pick_best_gguf(repo_files) except Exception as e: last_err = e - # 404 / RepoNotFound is permanent -- don't waste attempts. + # 404 / RepoNotFound is permanent -- don't retry err_name = type(e).__name__ if err_name in ( "RepositoryNotFoundError", @@ -1611,9 +1653,7 @@ def detect_gguf_model_remote( ) return cached - logger.warning( - f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}" - ) + logger.warning(f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}") return None @@ -1622,11 +1662,7 @@ def download_gguf_file( filename: str, hf_token: Optional[str] = None, ) -> str: - """ - Download a specific GGUF file from a HuggingFace repo. - - Returns the local path to the downloaded file. - """ + """Download a specific GGUF file from a HF repo; returns the local path.""" from huggingface_hub import hf_hub_download local_path = hf_hub_download( @@ -1637,35 +1673,29 @@ def download_gguf_file( return local_path -# Cache embedding detection results per session to avoid repeated HF API calls +# Cache embedding detection per session to avoid repeated HF API calls _embedding_detection_cache: Dict[tuple, bool] = {} def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: - """ - Detect embedding/sentence-transformer models using HuggingFace model metadata. + """Detect embedding/sentence-transformer models via HF metadata. - Uses a belt-and-suspenders approach combining three signals: - 1. "sentence-transformers" in model tags - 2. "feature-extraction" in model tags - 3. pipeline_tag is "sentence-similarity" or "feature-extraction" - - This catches all known embedding models including those like gte-modernbert - whose library_name is "transformers" rather than "sentence-transformers". + Combines three signals: "sentence-transformers" or "feature-extraction" in + tags, or pipeline_tag in {"sentence-similarity", "feature-extraction"}. + Catches models like gte-modernbert whose library_name is "transformers". Args: model_name: Model identifier (HF repo or local path) - hf_token: Optional HF token for accessing gated/private models + hf_token: Optional HF token for gated/private models Returns: - True if the model is an embedding model, False otherwise. - Defaults to False for local paths or on errors. + True if embedding model, else False (default for local paths or errors). """ cache_key = (model_name, hf_token) if cache_key in _embedding_detection_cache: return _embedding_detection_cache[cache_key] - # Local paths: check for sentence-transformer marker file (modules.json) + # Local paths: check for sentence-transformer marker (modules.json) if is_local_path(model_name): local_dir = normalize_path(model_name) is_emb = os.path.isfile(os.path.join(local_dir, "modules.json")) @@ -1746,15 +1776,12 @@ def _looks_like_lora_adapter(model_dir: Path) -> bool: ) -def scan_trained_models( - outputs_dir: str = str(outputs_root()), -) -> List[Tuple[str, str, str]]: - """ - Scan outputs folder for trained Studio models. +def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str, str]]: + """Scan outputs folder for trained Studio models. Returns: - List of tuples: [(display_name, model_path, model_type), ...] - model_type is "lora" for adapter runs and "merged" for full finetunes. + List of (display_name, model_path, model_type), where model_type is + "lora" for adapter runs or "merged" for full finetunes. """ trained_models = [] outputs_path = resolve_output_dir(outputs_dir) @@ -1775,7 +1802,7 @@ def scan_trained_models( trained_models.append((display_name, model_path, model_type)) logger.debug("Found trained model: %s (%s)", display_name, model_type) - # Sort by modification time (newest first) + # Sort by mtime, newest first trained_models.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True) logger.info( @@ -1793,16 +1820,14 @@ def scan_trained_models( def scan_exported_models( exports_dir: str = str(exports_root()), ) -> List[Tuple[str, str, str, Optional[str]]]: - """ - Scan exports folder for exported models (merged, LoRA, GGUF). + """Scan exports folder for exported models (merged, LoRA, GGUF). - Supports two directory layouts: - - Two-level: {run}/{checkpoint}/ (merged & LoRA exports) - - Flat: {name}-finetune-gguf/ (GGUF exports) + Supports two layouts: two-level {run}/{checkpoint}/ (merged & LoRA) and + flat {name}-finetune-gguf/ (GGUF). Returns: - List of tuples: [(display_name, model_path, export_type, base_model), ...] - export_type: "lora" | "merged" | "gguf" + List of (display_name, model_path, export_type, base_model), where + export_type is "lora" | "merged" | "gguf". """ results = [] exports_path = resolve_export_dir(exports_dir) @@ -1815,11 +1840,9 @@ def scan_exported_models( if not run_dir.is_dir(): continue - # Check for flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/) - # Filter out mmproj (vision projection) files — they aren't loadable as main models - gguf_files = [ - f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name) - ] + # Flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/). + # Skip mmproj (vision projection) files — not loadable as main models. + gguf_files = [f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name)] if gguf_files: base_model = None export_meta = run_dir / "export_metadata.json" @@ -1831,7 +1854,7 @@ def scan_exported_models( pass display_name = run_dir.name - model_path = str(gguf_files[0]) # path to the .gguf file + model_path = str(gguf_files[0]) results.append((display_name, model_path, "gguf", base_model)) logger.debug(f"Found GGUF export: {display_name}") continue @@ -1870,8 +1893,8 @@ def scan_exported_models( elif has_gguf: export_type = "gguf" gguf_list = list(_iter_gguf_files(checkpoint_dir)) - # Check checkpoint_dir first, then fall back to parent run_dir - # (export.py writes metadata to the top-level export directory) + # checkpoint_dir first, then run_dir (export.py writes + # metadata to the top-level export dir) for meta_dir in (checkpoint_dir, run_dir): export_meta = meta_dir / "export_metadata.json" try: @@ -1891,12 +1914,9 @@ def scan_exported_models( else: continue - # Fallback: read base model from the original training run's - # adapter_config.json in ./outputs/{run_name}/ + # Fallback: base model from ./outputs/{run_name}/adapter_config.json if not base_model: - outputs_adapter_cfg = ( - resolve_output_dir(run_dir.name) / "adapter_config.json" - ) + outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json" try: if outputs_adapter_cfg.exists(): cfg = json.loads(outputs_adapter_cfg.read_text()) @@ -1929,9 +1949,7 @@ def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]: config = json.load(f) base_model = config.get("base_model_name_or_path") if base_model: - logger.info( - "Detected base model from adapter_config.json: %s", base_model - ) + logger.info("Detected base model from adapter_config.json: %s", base_model) return base_model config_path = checkpoint_path_obj / "config.json" @@ -1982,31 +2000,21 @@ def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]: def get_base_model_from_lora(lora_path: str) -> Optional[str]: - """ - Read the base model name from a LoRA adapter's config. - - Args: - lora_path: Path to the LoRA adapter directory - - Returns: - Base model identifier or None if not found - """ + """Read the base model name from a LoRA adapter's config, or None.""" try: lora_path_obj = Path(lora_path) if not _looks_like_lora_adapter(lora_path_obj): return None - # Try adapter_config.json first + # adapter_config.json first adapter_config_path = lora_path_obj / "adapter_config.json" if adapter_config_path.exists(): with open(adapter_config_path, "r") as f: config = json.load(f) base_model = config.get("base_model_name_or_path") if base_model: - logger.info( - f"Detected base model from adapter_config.json: {base_model}" - ) + logger.info(f"Detected base model from adapter_config.json: {base_model}") return base_model # Fallback: try training_args.bin (requires torch) @@ -2026,13 +2034,10 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]: # except Exception as e: # logger.warning(f"Could not load training_args.bin: {e}") - # Last resort: parse from directory name - # Format: unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit_timestamp + # Last resort: parse from dir name (unsloth__) dir_name = lora_path_obj.name if dir_name.startswith("unsloth_"): - # Remove timestamp suffix (usually _1234567890) parts = dir_name.split("_") - # Reconstruct model name if len(parts) >= 2: model_parts = parts[1:-1] # Skip "unsloth" and timestamp base_model = "unsloth/" + "_".join(model_parts) @@ -2052,41 +2057,29 @@ UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "] def load_model_defaults(model_name: str) -> Dict[str, Any]: - """ - Load default training parameters for a model from YAML file. + """Load default training parameters for a model from a YAML file. - Args: - model_name: Model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-bnb-4bit") - - Returns: - Dictionary with default parameters from YAML file, or empty dict if not found - - The function looks for a YAML file in configs/model_defaults/ (including subfolders) - based on the model name or its aliases from MODEL_NAME_MAPPING. - If no specific file exists, it falls back to default.yaml. + Looks in configs/model_defaults/ (incl. subfolders) by model name or its + MODEL_NAME_MAPPING aliases, else falls back to default.yaml. Returns the + parameter dict, or {} if none found. """ try: - # Get the script directory to locate configs script_dir = Path(__file__).parent.parent.parent defaults_dir = script_dir / "assets" / "configs" / "model_defaults" - # First, check if model is in the mapping + # Check the mapping first if model_name.lower() in _REVERSE_MODEL_MAPPING: canonical_file = _REVERSE_MODEL_MAPPING[model_name.lower()] - # Search in subfolders and root for config_path in defaults_dir.rglob(canonical_file): if config_path.is_file(): with open(config_path, "r", encoding = "utf-8") as f: config = yaml.safe_load(f) or {} - logger.info( - f"Loaded model defaults from {config_path} (via mapping)" - ) + logger.info(f"Loaded model defaults from {config_path} (via mapping)") return config - # If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from - # adapter_config.json, or C:\Users\...\model on Windows), try matching - # the last 1-2 path components against the registry - # (e.g. "Spark-TTS-0.5B/LLM"). + # For local paths (e.g. /home/.../Spark-TTS-0.5B/LLM from + # adapter_config.json, or C:\Users\...\model on Windows), match the + # last 1-2 path components against the registry (e.g. "Spark-TTS-0.5B/LLM"). _is_local_path = is_local_path(model_name) # Normalize Windows backslash paths so Path().parts splits correctly # on POSIX/WSL hosts (pathlib treats backslashes as literals on Linux). @@ -2107,13 +2100,13 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]: ) return config - # Try exact model name match (for backward compatibility). - # For local filesystem paths, use only the directory basename to - # avoid passing absolute paths (e.g. C:\...) into rglob which - # raises "Non-relative patterns are unsupported" on Windows. + # Exact model name match (backward compatibility). For local paths, + # use only the dir basename to avoid passing absolute paths (e.g. + # C:\...) into rglob, which raises "Non-relative patterns are + # unsupported" on Windows. _lookup_name = Path(_normalized).name if _is_local_path else model_name model_filename = _lookup_name.replace("/", "_") + ".yaml" - # Search in subfolders and root + # Search subfolders and root for config_path in defaults_dir.rglob(model_filename): if config_path.is_file(): with open(config_path, "r", encoding = "utf-8") as f: @@ -2139,25 +2132,22 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]: @dataclass class ModelConfig: - """Configuration for a model to load""" + """Configuration for a model to load.""" identifier: str # Clean model identifier (org/name or path) display_name: str # Original UI display name path: str # Normalized filesystem path - is_local: bool # Is this a local file vs HF model? - is_cached: bool # Is this already in HF cache? - is_vision: bool # Is this a vision model? - is_lora: bool # Is this a lora adapter? - is_gguf: bool = False # Is this a GGUF model? - is_audio: bool = False # Is this a TTS audio model? - audio_type: Optional[str] = ( - None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac' - ) + is_local: bool # Local file vs HF model? + is_cached: bool # Already in HF cache? + is_vision: bool # Vision model? + is_lora: bool # LoRA adapter? + is_gguf: bool = False # GGUF model? + is_audio: bool = False # TTS audio model? + audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac' has_audio_input: bool = False # Accepts audio input (ASR/speech understanding) gguf_file: Optional[str] = None # Full path to the .gguf file (local mode) - gguf_mmproj_file: Optional[str] = ( - None # Full path to the mmproj .gguf file (vision projection) - ) + gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection) + gguf_mtp_file: Optional[str] = None # Full path to the separate MTP drafter (local mode) gguf_hf_repo: Optional[str] = ( None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF") ) @@ -2166,19 +2156,16 @@ class ModelConfig: @classmethod def from_lora_path( - cls, lora_path: str, hf_token: Optional[str] = None + cls, + lora_path: str, + hf_token: Optional[str] = None, ) -> Optional["ModelConfig"]: - """ - Create ModelConfig from a local LoRA adapter path. - - Automatically detects the base model from adapter config. + """Create ModelConfig from a local LoRA adapter path, auto-detecting the + base model from adapter config. Args: - lora_path: Path to LoRA adapter (e.g., "./outputs/unsloth_Meta-Llama-3.1_.../") + lora_path: Path to the LoRA adapter directory hf_token: HF token for vision detection - - Returns: - ModelConfig for the LoRA adapter """ try: lora_path_obj = Path(lora_path) @@ -2187,27 +2174,23 @@ class ModelConfig: logger.error(f"LoRA path does not exist: {lora_path}") return None - # Get base model base_model = get_base_model_from_lora(lora_path) if not base_model: logger.error(f"Could not determine base model for LoRA: {lora_path}") return None - # Check if base model is vision is_vision = is_vision_model(base_model, hf_token = hf_token) - - # Check if base model is audio audio_type = detect_audio_type(base_model, hf_token = hf_token) display_name = lora_path_obj.name - identifier = lora_path # Use path as identifier for local LoRAs + identifier = lora_path # path is the identifier for local LoRAs return cls( identifier = identifier, display_name = display_name, path = lora_path, is_local = True, - is_cached = True, # Local LoRAs are always "cached" + is_cached = True, # local LoRAs are always cached is_vision = is_vision, is_lora = True, is_audio = audio_type is not None and audio_type != "audio_vlm", @@ -2228,25 +2211,18 @@ class ModelConfig: is_lora: bool = False, gguf_variant: Optional[str] = None, ) -> Optional["ModelConfig"]: - """ - Create ModelConfig from a clean model identifier. - - For FastAPI routes where the frontend sends sanitized model paths. - No Gradio dropdown parsing - expects clean identifiers like: - - "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit" - - "./outputs/my_lora_adapter" - - "/absolute/path/to/model" + """Create ModelConfig from a clean model identifier (HF repo or local + path), for FastAPI routes that send sanitized paths. Args: model_id: Clean model identifier (HF repo name or local path) hf_token: Optional HF token for vision detection on gated models is_lora: Whether this is a LoRA adapter - gguf_variant: Optional GGUF quantization variant (e.g. "Q4_K_M"). - For remote GGUF repos, specifies which quant to load via -hf. - If None, auto-selects using _pick_best_gguf(). + gguf_variant: Optional GGUF quant variant (e.g. "Q4_K_M") to load + via -hf for remote repos; None auto-selects via _pick_best_gguf(). Returns: - ModelConfig or None if configuration cannot be created + ModelConfig or None if it cannot be created. """ if not model_id or not model_id.strip(): return None @@ -2260,8 +2236,8 @@ class ModelConfig: identifier = f"unsloth/{identifier}" path = identifier - # Preserve requested casing, but if a case-variant already exists in local HF cache, - # reuse that exact repo_id spelling to avoid one-time re-downloads after #2592. + # Reuse a cached case-variant's exact repo_id spelling to avoid + # one-time re-downloads after #2592. if not is_local: resolved_identifier = resolve_cached_repo_id_case(identifier) if resolved_identifier != identifier: @@ -2283,12 +2259,12 @@ class ModelConfig: display_name = Path(gguf_file).stem logger.info(f"Detected local GGUF model: {gguf_file}") - # Detect vision: check if base model is vision, then look for mmproj + # Vision: check base model, then look for mmproj mmproj_file = None gguf_is_vision = False gguf_dir = Path(gguf_file).parent - # Determine if this is a vision model from export metadata + # Is this a vision model, per export metadata? base_is_vision = False meta_path = gguf_dir / "export_metadata.json" if meta_path.exists(): @@ -2301,23 +2277,20 @@ class ModelConfig: except Exception as e: logger.debug(f"Could not read export metadata: {e}") - # If vision (or mmproj happens to exist), find the mmproj - # file. The recursive variant scan in - # ``_find_local_gguf_by_variant`` may have returned a - # weight file inside a quant-named subdir (e.g. - # ``.../BF16/foo.gguf``) while ``mmproj-*.gguf`` lives - # at the snapshot root. Pass ``search_root=path`` so - # ``detect_mmproj_file`` walks up to the snapshot root - # instead of seeing only the weight file's immediate - # parent. + # Pass search_root=path so detect_mmproj_file walks up to the + # snapshot root: the weight may sit in a quant subdir while + # mmproj-*.gguf lives at the root. mmproj_file = detect_mmproj_file(gguf_file, search_root = path) if mmproj_file: gguf_is_vision = True logger.info(f"Detected mmproj for vision: {mmproj_file}") elif base_is_vision: - logger.warning( - f"Base model is vision but no mmproj file found in {gguf_dir}" - ) + logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}") + + # Separate MTP drafter sibling (Gemma 4), mirroring mmproj. + mtp_file = detect_mtp_file(gguf_file, search_root = path) + if mtp_file: + logger.info(f"Detected MTP drafter: {mtp_file}") return cls( identifier = identifier, @@ -2330,13 +2303,13 @@ class ModelConfig: is_gguf = True, gguf_file = gguf_file, gguf_mmproj_file = mmproj_file, + gguf_mtp_file = mtp_file, ) else: - # Check if the HF repo contains GGUF files + # Does the HF repo contain GGUF files? gguf_filename = detect_gguf_model_remote(identifier, hf_token = hf_token) if gguf_filename: - # Preflight: verify llama-server binary exists BEFORE user waits - # for a multi-GB download that llama-server handles natively + # Preflight: verify llama-server binary exists before a multi-GB download from core.inference.llama_cpp import LlamaCppBackend if not LlamaCppBackend._find_llama_server_binary(): @@ -2345,11 +2318,10 @@ class ModelConfig: "Run setup.sh to build it, or set LLAMA_SERVER_PATH." ) - # Use list_gguf_variants() to detect vision & resolve variant + # list_gguf_variants() detects vision & resolves the variant variants, has_vision = list_gguf_variants(identifier, hf_token = hf_token) variant = gguf_variant - if not variant: - # Auto-select best quantization + if not variant: # auto-select best quant variant_filenames = [v.filename for v in variants] best = _pick_best_gguf(variant_filenames) if best: @@ -2376,18 +2348,14 @@ class ModelConfig: gguf_variant = variant, ) - # Auto-detect LoRA for local paths (check adapter_config.json on disk) + # Auto-detect LoRA for local paths (adapter_config.json on disk) if not is_lora and is_local: detected_base = ( - get_base_model_from_lora(path) - if _looks_like_lora_adapter(Path(path)) - else None + get_base_model_from_lora(path) if _looks_like_lora_adapter(Path(path)) else None ) if detected_base: is_lora = True - logger.info( - f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})" - ) + logger.info(f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})") # Auto-detect LoRA for remote HF models. When offline, huggingface_hub # raises OfflineModeIsEnabled in ~0ms; we fall through to the cache. @@ -2401,18 +2369,14 @@ class ModelConfig: is_lora = True logger.info(f"Auto-detected remote LoRA adapter: '{identifier}'") except Exception as e: - logger.debug( - f"Could not check remote LoRA status for '{identifier}': {e}" - ) + logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}") - # API may have failed; adapter_config.json may still be cached. + # API may have failed; adapter_config.json could still be cached. if not is_lora: for snap in _iter_hf_cache_snapshots(identifier): if (snap / "adapter_config.json").is_file(): is_lora = True - logger.info( - f"Auto-detected cached LoRA adapter: '{identifier}'" - ) + logger.info(f"Auto-detected cached LoRA adapter: '{identifier}'") break # Handle LoRA adapters @@ -2422,13 +2386,11 @@ class ModelConfig: # Local LoRA: read adapter_config.json from disk base_model = get_base_model_from_lora(path) else: - # Remote LoRA: download adapter_config.json from HF + # Remote LoRA: fetch adapter_config.json from HF try: from huggingface_hub import hf_hub_download - config_path = hf_hub_download( - identifier, "adapter_config.json", token = hf_token - ) + config_path = hf_hub_download(identifier, "adapter_config.json", token = hf_token) with open(config_path, "r") as f: adapter_config = json.load(f) base_model = adapter_config.get("base_model_name_or_path") @@ -2475,10 +2437,7 @@ class ModelConfig: hf_token: Optional[str] = None, is_lora: bool = False, ) -> Optional["ModelConfig"]: - """ - Create a universal ModelConfig from UI dropdown/search selections. - Handles base models and LoRA adapters. - """ + """Create a ModelConfig from UI dropdown/search selections (base models and LoRAs).""" selected = None if search_value and search_value.strip(): selected = search_value.strip() @@ -2490,18 +2449,16 @@ class ModelConfig: display_name = selected - # Use the correct 'local_models' parameter to resolve display names + # Resolve display names via the 'local_models' parameter if " (Active)" in selected or " (Ready)" in selected: - clean_display_name = selected.replace(" (Active)", "").replace( - " (Ready)", "" - ) + clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "") if local_models: for local_display, local_path in local_models: if local_display == clean_display_name: selected = local_path break - # Clean all UI status indicators to get the final identifier + # Strip all UI status indicators to get the final identifier identifier = selected for status in UI_STATUS_INDICATORS: identifier = identifier.replace(status, "") @@ -2521,23 +2478,23 @@ class ModelConfig: identifier = resolved_identifier path = resolved_identifier - # --- Logic for Base Model and Vision Detection --- + # --- Base Model and Vision Detection --- base_model = None is_vision = False if is_lora: - # For a LoRA, we MUST find its base model. + # A LoRA MUST have a base model. base_model = get_base_model_from_lora(path) if not base_model: logger.warning( f"Could not determine base model for LoRA '{path}'. Cannot create config." ) - return None # Cannot proceed without a base model + return None # cannot proceed without a base model - # A LoRA's vision capability is determined by its base model. + # A LoRA's vision capability comes from its base model. is_vision = is_vision_model(base_model, hf_token = hf_token) else: - # For a base model, just check its own vision status. + # Base model: check its own vision status. is_vision = is_vision_model(identifier, hf_token = hf_token) from utils.paths import is_model_cached @@ -2552,5 +2509,5 @@ class ModelConfig: is_cached = is_cached, is_vision = is_vision, is_lora = is_lora, - base_model = base_model, # This will be None for base models, and populated for LoRAs + base_model = base_model, # None for base models, set for LoRAs ) diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py index a69dfab532..7d8514abc8 100644 --- a/studio/backend/utils/native_path_leases.py +++ b/studio/backend/utils/native_path_leases.py @@ -68,9 +68,7 @@ def native_path_leases_supported() -> bool: return True -def child_env_without_native_path_secret( - env: Mapping[str, str] | None = None, -) -> dict[str, str]: +def child_env_without_native_path_secret(env: Mapping[str, str] | None = None) -> dict[str, str]: """Return a child-process env with the native path lease secret removed.""" if env is None: @@ -82,11 +80,7 @@ def child_env_without_native_path_secret( return cleaned -def run_without_native_path_secret( - target: Callable[..., Any], - *args: Any, - **kwargs: Any, -) -> Any: +def run_without_native_path_secret(target: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: """Run a multiprocessing child target without the native path lease secret.""" global _CACHED_LEASE_SECRET, _SCRUB_SAVED_SECRET @@ -153,9 +147,7 @@ def verify_native_path_lease( raise NativePathLeaseError("Native path is no longer accessible.") from exc _reject_network_or_device_path(resolved) if not _same_native_path(resolved, path): - raise NativePathLeaseError( - "Native path grant no longer resolves to the selected path." - ) + raise NativePathLeaseError("Native path grant no longer resolves to the selected path.") grant = NativePathGrant( operation = str(payload["operation"]), @@ -219,9 +211,7 @@ def _decode_secret() -> bytes: if encoded is None and _SCRUB_SAVED_SECRET is not None: encoded = _SCRUB_SAVED_SECRET if not encoded: - raise NativePathLeaseError( - "Native path grants require the managed desktop backend." - ) + raise NativePathLeaseError("Native path grants require the managed desktop backend.") try: secret = _b64decode(encoded) except Exception as exc: @@ -272,9 +262,7 @@ def _validate_payload( ) missing = [key for key in required if key not in payload] if missing: - raise NativePathLeaseError( - "Native path grant payload is missing required fields." - ) + raise NativePathLeaseError("Native path grant payload is missing required fields.") if _required_int(payload, "version") != 1: raise NativePathLeaseError("Native path grant version is unsupported.") if payload["operation"] != operation: @@ -353,19 +341,13 @@ def _reject_network_or_device_path(path: Path) -> None: rest = normalized[4:] is_local_drive = len(rest) >= 3 and rest[0].isalpha() and rest[1:3] == ":\\" if not is_local_drive: - raise NativePathLeaseError( - "Network paths are not supported for native grants." - ) + raise NativePathLeaseError("Network paths are not supported for native grants.") elif normalized.startswith("\\\\"): - raise NativePathLeaseError( - "Network paths are not supported for native grants." - ) + raise NativePathLeaseError("Network paths are not supported for native grants.") if os.name != "nt": for root in ("/dev", "/proc", "/sys"): if path.is_relative_to(root): - raise NativePathLeaseError( - "Device and virtual filesystem paths are not supported." - ) + raise NativePathLeaseError("Device and virtual filesystem paths are not supported.") if "\x00" in text: raise NativePathLeaseError("Native path contains invalid characters.") @@ -397,9 +379,7 @@ def _optional_int(value: Any) -> int | None: def _required_int(payload: dict[str, Any], key: str) -> int: raw = payload.get(key) if raw is None: - raise NativePathLeaseError( - "Native path grant payload is missing required fields." - ) + raise NativePathLeaseError("Native path grant payload is missing required fields.") try: return int(raw) except (TypeError, ValueError) as exc: diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 92191dccdd..6913a3be73 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Path utilities for model and dataset handling -""" +"""Path utilities for model and dataset handling.""" from .path_utils import ( normalize_path, @@ -25,6 +23,11 @@ from .storage_roots import ( auth_root, auth_db_path, studio_db_path, + rag_root, + rag_db_path, + rag_uploads_root, + documents_root, + project_workspaces_root, tmp_root, seed_uploads_root, unstructured_seed_cache_root, @@ -44,6 +47,10 @@ from .storage_roots import ( resolve_dataset_path, ) +# Re-export shim: mark project-path helpers as used so the import-hoist +# safety net does not flag them as unused. +_REEXPORTED = (documents_root, project_workspaces_root) + __all__ = [ "normalize_path", "is_local_path", @@ -62,6 +69,11 @@ __all__ = [ "auth_root", "auth_db_path", "studio_db_path", + "rag_root", + "rag_db_path", + "rag_uploads_root", + "documents_root", + "project_workspaces_root", "tmp_root", "seed_uploads_root", "unstructured_seed_cache_root", @@ -80,3 +92,6 @@ __all__ = [ "resolve_tensorboard_dir", "resolve_dataset_path", ] + +# Bind the re-exports so the import-hoist verifier counts them as used. +_ = (rag_root, rag_db_path, rag_uploads_root) diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py index 9ef9a2dd92..e8dabc8954 100644 --- a/studio/backend/utils/paths/path_utils.py +++ b/studio/backend/utils/paths/path_utils.py @@ -17,7 +17,7 @@ logger = get_logger(__name__) # Per-process cache to avoid repeated cache-dir scans for the same identifier. _CACHE_CASE_RESOLUTION_MEMO: dict[str, str] = {} -# Lightweight instrumentation counters for operational visibility. +# Instrumentation counters for operational visibility. _CACHE_CASE_RESOLUTION_STATS: dict[str, int] = { "calls": 0, "memo_hits": 0, @@ -44,27 +44,17 @@ _IS_WSL: bool = _is_wsl() def normalize_path(path: str) -> str: - """ - Normalize filesystem paths for cross-platform use. + """Normalize filesystem paths for cross-platform use. - On WSL, converts Windows drive-letter paths to ``/mnt//...``. - On native Windows, keeps the drive letter and normalizes separators. - On Linux/macOS (non-WSL), paths are returned with forward slashes. - - Examples (WSL): - C:\\Users\\... -> /mnt/c/Users/... - Examples (native Windows): - C:\\Users\\... -> C:/Users/... - Examples (Linux/macOS): - /home/user/... -> /home/user/... (unchanged) + WSL maps drive-letter paths to ``/mnt//...``; native Windows keeps + the drive and normalizes separators; elsewhere slashes are forward-only. """ if not path: return path # Handle Windows drive letters (C:\\ or c:\\) if len(path) >= 3 and path[1] == ":" and path[2] in ("\\", "/"): - # Only map to /mnt// when running under WSL; - # on native Windows the drive letter must be preserved. + # Map to /mnt// only under WSL; native Windows keeps the drive letter. if _IS_WSL: drive = path[0].lower() rest = path[3:].replace("\\", "/") @@ -86,7 +76,7 @@ def is_local_path(path: str) -> bool: if not path: return False - # If it exists on disk, treat as local (covers relative paths like "outputs/foo"). + # Exists on disk → local (covers relative paths like "outputs/foo"). try: if Path(normalize_path(path)).expanduser().exists(): return True @@ -122,7 +112,7 @@ def is_model_cached(model_name: str) -> bool: if not cache_path: return False - # Check for actual model files + # Check for model files for suffix in [".safetensors", ".bin", ".json"]: if list(cache_path.rglob(f"*{suffix}")): return True @@ -134,7 +124,6 @@ def _hf_hub_cache_dir() -> Path: """Return HF cache root honoring HF_HUB_CACHE when available.""" try: from huggingface_hub.constants import HF_HUB_CACHE - return Path(HF_HUB_CACHE) except Exception as exc: logger.debug( @@ -147,9 +136,9 @@ def _hf_hub_cache_dir() -> Path: def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str: """Resolve repo_id to the exact casing already present in local HF cache. - Policy: prefer the requested/canonical repo_id, but if a case-variant already - exists in local HF cache, reuse that exact cached spelling. This avoids - duplicate downloads while preserving user intent whenever possible. + Policy: prefer the requested/canonical repo_id, but reuse a case-variant's + exact cached spelling if one already exists in local HF cache. Avoids + duplicate downloads while preserving user intent where possible. """ _CACHE_CASE_RESOLUTION_STATS["calls"] += 1 @@ -164,8 +153,7 @@ def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str: expected_dir = f"models--{model_name.replace('/', '--')}" - # Always check the exact-case path first so a newly-appeared exact match - # wins over any previously memoized variant. + # Exact-case path first so a new exact match beats a memoized variant. exact_path = cache_dir / expected_dir if exact_path.is_dir(): if use_memo: @@ -173,8 +161,7 @@ def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str: _CACHE_CASE_RESOLUTION_STATS["exact_hits"] += 1 return model_name - # Validate memoized entries still exist on disk before returning them. - # This prevents stale results when cache dirs are deleted/recreated. + # Revalidate memoized entries on disk to avoid stale results. if use_memo: cached = _CACHE_CASE_RESOLUTION_MEMO.get(model_name) if cached is not None: @@ -182,7 +169,7 @@ def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str: if cached_path.is_dir(): _CACHE_CASE_RESOLUTION_STATS["memo_hits"] += 1 return cached - # Stale entry -- drop it and re-scan below. + # Stale entry -- drop it and re-scan below _CACHE_CASE_RESOLUTION_MEMO.pop(model_name, None) expected_lower = expected_dir.lower() @@ -201,7 +188,7 @@ def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str: candidates.append(repo_part.replace("--", "/")) if candidates: - # Deterministic tie-break if multiple case variants coexist. + # Deterministic tie-break if multiple case variants coexist resolved = sorted(candidates)[0] if len(candidates) > 1: _CACHE_CASE_RESOLUTION_STATS["tie_breaks"] += 1 diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 763d18bf3e..2d4f5ac243 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -11,9 +11,9 @@ import tempfile def _infer_studio_home_from_venv() -> Path | None: - """Return parent dir of sys.prefix as STUDIO_HOME if running from an + """Return parent of sys.prefix as STUDIO_HOME when running from an installer-managed unsloth_studio venv. Sentinel-gated (share/studio.conf - or bin shim) so a developer venv named unsloth_studio is not misidentified. + or bin shim) so a dev venv named unsloth_studio isn't misidentified. """ try: prefix = Path(sys.prefix).resolve() @@ -38,8 +38,8 @@ def studio_root() -> Path: """Studio install root. Priority: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then sys.prefix - inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins when - both are set (the more specific signal beats the generic alias). + inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins if + both are set (specific signal beats generic alias). """ override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip() if not override: @@ -56,10 +56,15 @@ def studio_root() -> Path: def cache_root() -> Path: - """Central cache directory for all studio downloads (models, datasets, etc.).""" + """Central cache dir for all studio downloads (models, datasets, etc.).""" return studio_root() / "cache" +def studio_bin_root() -> Path: + """Dir for Studio-managed executables (the `unsloth` shim, downloaded tools like cloudflared).""" + return studio_root() / "bin" + + def assets_root() -> Path: return studio_root() / "assets" @@ -96,6 +101,53 @@ def studio_db_path() -> Path: return studio_root() / "studio.db" +def rag_root() -> Path: + """Root directory for retrieval-augmented-generation state (db + uploads).""" + return studio_root() / "rag" + + +def rag_db_path() -> Path: + """SQLite file holding RAG documents, chunks, FTS5 + sqlite-vec indexes.""" + return rag_root() / "rag.db" + + +def rag_uploads_root() -> Path: + """Directory where uploaded source documents are stored for ingestion.""" + return rag_root() / "uploads" + + +def _xdg_user_dir(key: str) -> Path | None: + config = Path.home() / ".config" / "user-dirs.dirs" + try: + lines = config.read_text(encoding = "utf-8").splitlines() + except OSError: + return None + prefix = f"{key}=" + for line in lines: + line = line.strip() + if not line.startswith(prefix): + continue + value = line[len(prefix) :].strip().strip('"') + if not value: + return None + return Path(value.replace("$HOME", str(Path.home()))).expanduser() + return None + + +def documents_root() -> Path: + override = (os.environ.get("UNSLOTH_STUDIO_DOCUMENTS_HOME") or "").strip() + if override: + return Path(override).expanduser() + return _xdg_user_dir("XDG_DOCUMENTS_DIR") or (Path.home() / "Documents") + + +def project_workspaces_root() -> Path: + override = (os.environ.get("UNSLOTH_STUDIO_PROJECTS_HOME") or "").strip() + if override: + return Path(override).expanduser() + return documents_root() / "Unsloth Studio" / "Projects" + + def tmp_root() -> Path: return Path(tempfile.gettempdir()) / "unsloth-studio" @@ -126,16 +178,15 @@ def ensure_dir(path: Path) -> Path: def legacy_hf_cache_dir() -> Path: - """Old Unsloth-specific HF hub cache, kept for backward-compat scanning.""" + """Old Unsloth-specific HF hub cache, kept for backward-compat scans.""" return cache_root() / "huggingface" / "hub" def hf_default_cache_dir() -> Path: - """Return the platform default HuggingFace hub cache (ignoring env overrides). + """Platform default HuggingFace hub cache (ignoring env overrides). - This is the location HF uses when no ``HF_HUB_CACHE`` / ``HF_HOME`` - env var is set. We scan it so that models a user downloaded *before* - installing Unsloth Studio are still discovered. + Where HF caches when no ``HF_HUB_CACHE`` / ``HF_HOME`` is set. Scanned + so models downloaded *before* installing Unsloth Studio are discovered. """ return Path.home() / ".cache" / "huggingface" / "hub" @@ -151,7 +202,7 @@ def lmstudio_model_dirs() -> list[Path]: seen.add(resolved) dirs.append(p) - # 1. Check LM Studio settings.json for custom downloads folder + # LM Studio settings.json custom downloads folder settings_path = Path.home() / ".lmstudio" / "settings.json" if settings_path.is_file(): try: @@ -163,10 +214,10 @@ def lmstudio_model_dirs() -> list[Path]: except Exception: pass - # 2. LM Studio current default models directory (all platforms) + # LM Studio default models directory (all platforms) _add(Path.home() / ".lmstudio" / "models") - # 3. Legacy LM Studio cache location + # Legacy LM Studio cache location _add(Path.home() / ".cache" / "lm-studio" / "models") return dirs @@ -175,17 +226,17 @@ def lmstudio_model_dirs() -> list[Path]: def well_known_model_dirs() -> list[Path]: """Return directories commonly used by other local LLM tools. - Used by the folder browser to offer quick-pick chips. Returns only - paths that exist on disk, so the UI never shows dead chips. Order - reflects a rough "likelihood the user has models here" -- LM Studio - and Ollama first, then the generic fallbacks. + Backs the folder browser's quick-pick chips. Returns only paths that + exist on disk, so the UI never shows dead chips. Order reflects rough + likelihood of models being there -- LM Studio and Ollama first, then + generic fallbacks. """ candidates: list[Path] = [] # LM Studio (reuses the logic above, including settings.json override) candidates.extend(lmstudio_model_dirs()) - # Ollama -- both the user-level and common system-wide install paths + # Ollama -- user-level and common system-wide install paths # (https://github.com/ollama/ollama/issues/733). ollama_env = os.environ.get("OLLAMA_MODELS") if ollama_env: @@ -197,11 +248,11 @@ def well_known_model_dirs() -> list[Path]: # HF hub cache root (separate from the explicit HF cache chip) candidates.append(Path.home() / ".cache" / "huggingface" / "hub") - # Generic "my models" spots users tend to drop things into + # Generic "my models" spots users drop things into for name in ("models", "Models"): candidates.append(Path.home() / name) - # Deduplicate while preserving order; keep only extant dirs + # Dedupe preserving order; keep only extant dirs out: list[Path] = [] seen: set[str] = set() for p in candidates: @@ -218,22 +269,14 @@ def well_known_model_dirs() -> list[Path]: def _setup_cache_env() -> None: - """Set cache environment variables for HuggingFace, uv, and vLLM. + """Set cache env vars for HuggingFace, uv, and vLLM. - Respects the standard HF cache resolution chain: explicit ``HF_HOME`` - / ``HF_HUB_CACHE`` env vars take priority, then ``XDG_CACHE_HOME``, - then the platform default (``~/.cache/huggingface``). The legacy - Unsloth cache is still *scanned* for models but is never set as the - active download target. - - Only sets variables that are not already set by the user, so - explicit overrides (e.g. HF_HOME=/data/hf) are respected. - Works on Linux, macOS, and Windows. + Respects the standard HF cache chain (explicit HF_HOME / HF_HUB_CACHE, + then XDG_CACHE_HOME, then ~/.cache/huggingface) and only sets vars the + user hasn't, so explicit overrides are honored. """ root = cache_root() - xdg_cache = Path( - os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache") - ).expanduser() + xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() hf_default = xdg_cache / "huggingface" defaults: dict[str, str] = { "HF_HOME": str(hf_default), @@ -266,9 +309,7 @@ def ensure_studio_directories() -> None: _setup_cache_env() -def _clean_relative_path( - path_value: str, *, strip_prefixes: tuple[str, ...] = () -) -> Path: +def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = ()) -> Path: path = Path(path_value).expanduser() parts = [part for part in path.parts if part not in ("", ".")] while parts and parts[0] in strip_prefixes: @@ -287,8 +328,7 @@ def _assert_contained(resolved: Path, root: Path) -> None: resolved_real.relative_to(root_real) except ValueError as exc: raise ValueError( - f"path escapes root: {resolved!s} -> {resolved_real!s} " - f"is not under {root_real!s}" + f"path escapes root: {resolved!s} -> {resolved_real!s} " f"is not under {root_real!s}" ) from exc @@ -300,8 +340,8 @@ def resolve_under_root( ) -> Path: """Resolve ``path_value`` and assert the result is under ``root``. - Absolutes are accepted only if already contained (so internal pre-resolved - paths re-enter idempotently); user-facing schemas reject absolutes upstream. + Absolutes are accepted only if already contained (so pre-resolved + internal paths re-enter idempotently); schemas reject absolutes upstream. """ if not path_value or not str(path_value).strip(): return root @@ -362,9 +402,7 @@ def resolve_dataset_path(path_value: str) -> Path: return path except ValueError: continue - raise ValueError( - f"dataset path must be relative or under a dataset root: {raw!r}" - ) + raise ValueError(f"dataset path must be relative or under a dataset root: {raw!r}") parts = [part for part in Path(path_value).parts if part not in ("", ".")] if parts[:2] == ["assets", "datasets"]: diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py index 70059f8a3c..9c18070fbb 100644 --- a/studio/backend/utils/studio_version.py +++ b/studio/backend/utils/studio_version.py @@ -15,6 +15,7 @@ _DEV_VERSION = "dev" _GIT_TIMEOUT_SECONDS = 1.0 _STUDIO_TAG_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$") _GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$") +_GIT_BRANCH_RE = re.compile(r"^[0-9A-Za-z._/-]+$") _MAX_VERSION_LENGTH = 64 @@ -39,9 +40,7 @@ def _path_is_in_site_packages(path: Path) -> bool: def _is_source_checkout(repo_root: Path) -> bool: - return (repo_root / ".git").exists() and not _path_is_in_site_packages( - Path(__file__).resolve() - ) + return (repo_root / ".git").exists() and not _path_is_in_site_packages(Path(__file__).resolve()) def _exact_git_studio_tag(repo_root: Path) -> str | None: @@ -73,17 +72,49 @@ def _exact_git_studio_tag(repo_root: Path) -> str | None: return tag if is_valid_studio_release_version(tag) else None +def _git_branch(repo_root: Path) -> str | None: + try: + result = subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], + cwd = repo_root, + check = False, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = _GIT_TIMEOUT_SECONDS, + ) + except (OSError, subprocess.TimeoutExpired): + return None + + if result.returncode != 0: + return None + + branch = result.stdout.strip() + # "HEAD" means detached, e.g. a tag or commit checkout. + if ( + not branch + or branch == "HEAD" + or len(branch) > _MAX_VERSION_LENGTH + or _GIT_BRANCH_RE.fullmatch(branch) is None + ): + return None + return branch + + def get_studio_version(repo_root: Path | None = None) -> str: """Return the installed Studio release tag for display, or ``dev``. - This value is intentionally separate from the PyPI ``unsloth`` package - version used by update checks. It never performs network requests. + Intentionally separate from the PyPI ``unsloth`` package version used by + update checks. Never performs network requests. """ resolved_repo_root = repo_root or _repo_root() if _is_source_checkout(resolved_repo_root): git_tag = _exact_git_studio_tag(resolved_repo_root) - return git_tag if git_tag is not None else _DEV_VERSION + if git_tag is not None: + return git_tag + branch = _git_branch(resolved_repo_root) + return f"GitHub {branch}" if branch is not None else _DEV_VERSION stamped_version = _studio_release_build.STUDIO_RELEASE_VERSION if is_valid_studio_release_version(stamped_version): diff --git a/studio/backend/utils/subprocess_compat.py b/studio/backend/utils/subprocess_compat.py index bedf8cf2e6..f2fa6eadc7 100644 --- a/studio/backend/utils/subprocess_compat.py +++ b/studio/backend/utils/subprocess_compat.py @@ -8,10 +8,9 @@ import sys def windows_hidden_subprocess_kwargs() -> dict[str, object]: - """Return Windows-only subprocess kwargs that suppress console windows. + """Windows-only subprocess kwargs that suppress console windows. - On non-Windows platforms returns an empty dict so callers can always - unpack the result into ``subprocess.run`` / ``subprocess.Popen`` via + Empty dict off Windows, so callers can always unpack via ``**windows_hidden_subprocess_kwargs()``. """ if sys.platform != "win32": diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index c23857e0a4..6e1571d5ce 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -1,17 +1,17 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Automatic transformers version switching. +"""Automatic transformers version switching. Some newer model architectures (Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B MoE, -tiny_qwen3_moe) require transformers>=5.3.0, while Gemma 4 models require -transformers>=5.5.0. Everything else needs the default 4.57.x that ships -with Unsloth. +tiny_qwen3_moe) require transformers>=5.3.0, while Gemma 4 models require a +newer 5.x sidecar. Everything else needs the default 4.57.x that ships with +Unsloth. Two separate target directories are maintained: - .venv_t5_530/ — transformers 5.3.0 (Ministral-3, GLM, Qwen3 MoE, etc.) - .venv_t5_550/ — transformers 5.5.0 (Gemma 4) + - .venv_t5_510/ — transformers 5.10.2 (Gemma 4 Unified / 12B) When loading a LoRA adapter with a custom name, we resolve the base model from ``adapter_config.json`` and check *that* against the model list. @@ -57,8 +57,7 @@ def _env_offline() -> bool: # Detection # --------------------------------------------------------------------------- -# Lowercase substrings — if ANY appears anywhere in the lowered model name, -# we need transformers 5.3.0. +# Lowercase substrings — any match in the lowered model name needs transformers 5.3.0. TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( "ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512 "glm-4.7-flash", # GLM-4.7-Flash @@ -69,13 +68,32 @@ TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( "lfm2.5-vl-450m", # LiquidAI/LFM2.5-VL-450M ) -# Lowercase substrings for models that require transformers 5.5.0 (checked first). +# Lowercase substrings for models that require transformers 5.10.x (checked first). +TRANSFORMERS_510_MODEL_SUBSTRINGS: tuple[str, ...] = ( + "gemma-4-12b", # Gemma 4 Unified 12B + "gemma4-12b", +) + +# Lowercase substrings for models that require the Gemma 4 transformers 5.5 sidecar. TRANSFORMERS_550_MODEL_SUBSTRINGS: tuple[str, ...] = ( "gemma-4", # Gemma-4 (E2B-it, E4B-it, 31B-it, 26B-A4B-it) "gemma4", # Gemma-4 alternate naming "qwen3.6", ) +# Architecture classes / model_type values that require transformers 5.10.x. +# Checked via config.json (local or HuggingFace). +_TRANSFORMERS_510_ARCHITECTURES: set[str] = { + "Gemma4UnifiedForConditionalGeneration", + "Gemma4AssistantForCausalLM", + "Gemma4UnifiedAssistantForCausalLM", +} +_TRANSFORMERS_510_MODEL_TYPES: set[str] = { + "gemma4_unified", + "gemma4_assistant", + "gemma4_unified_assistant", +} + # Architecture classes / model_type values that require transformers 5.5.0. # Checked via config.json (local or HuggingFace). _TRANSFORMERS_550_ARCHITECTURES: set[str] = { @@ -85,30 +103,35 @@ _TRANSFORMERS_550_MODEL_TYPES: set[str] = { "gemma4", } -# Tokenizer classes that only exist in transformers>=5.x +# Tokenizer classes that only exist in transformers>=5.x. _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = { "TokenizersBackend", } -# Cache for dynamic tokenizer_config.json lookups to avoid repeated fetches +# Cache for dynamic tokenizer_config.json lookups (avoids repeated fetches). _tokenizer_class_cache: dict[str, bool] = {} -# Cache for dynamic config.json lookups (architecture/model_type checks) +# Cache for dynamic config.json lookups (architecture/model_type checks). +_config_json_cache: dict[str, dict | None] = {} +_config_needs_510_cache: dict[str, bool] = {} _config_needs_550_cache: dict[str, bool] = {} # Versions +TRANSFORMERS_510_VERSION = "5.10.2" TRANSFORMERS_550_VERSION = "5.5.0" TRANSFORMERS_530_VERSION = "5.3.0" TRANSFORMERS_DEFAULT_VERSION = "4.57.6" -# Backwards-compat alias — points to 5.5.0 (the highest 5.x tier). -# Consumers should prefer TRANSFORMERS_530_VERSION / TRANSFORMERS_550_VERSION. -TRANSFORMERS_5_VERSION = TRANSFORMERS_550_VERSION +# Backwards-compat alias — points to the highest 5.x tier. +# Consumers should prefer TRANSFORMERS_510_VERSION / TRANSFORMERS_550_VERSION / +# TRANSFORMERS_530_VERSION. +TRANSFORMERS_5_VERSION = TRANSFORMERS_510_VERSION # Pre-installed directories — created by setup.sh / setup.ps1. from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402 _VENV_T5_530_DIR = str(_studio_root() / ".venv_t5_530") _VENV_T5_550_DIR = str(_studio_root() / ".venv_t5_550") +_VENV_T5_510_DIR = str(_studio_root() / ".venv_t5_510") # Backwards-compat alias _VENV_T5_DIR = _VENV_T5_550_DIR @@ -116,25 +139,42 @@ _VENV_T5_DIR = _VENV_T5_550_DIR def activate_transformers_for_subprocess(model_name: str) -> None: """Activate the correct transformers version in a subprocess worker. - Call this BEFORE any ML imports. Resolves LoRA adapters to their base - model, determines the required tier, and prepends the appropriate - ``.venv_t5_*`` directory to ``sys.path``. Also propagates the path - via ``PYTHONPATH`` for child processes (e.g. GGUF converter). - - Used by training, inference, and export workers. + Call BEFORE any ML imports. Resolves LoRA adapters to their base model, + determines the required tier, prepends the appropriate ``.venv_t5_*`` dir to + ``sys.path``, and propagates it via ``PYTHONPATH`` for child processes + (e.g. GGUF converter). Used by training, inference, and export workers. """ resolved = _resolve_base_model(model_name) tier = get_transformers_tier(resolved) - if tier == "550": + if tier == "510": + if not _ensure_venv_t5_510_exists(): + raise RuntimeError( + f"Cannot activate transformers {TRANSFORMERS_510_VERSION}: " + f".venv_t5_510 missing at {_VENV_T5_510_DIR}" + ) + if _VENV_T5_510_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_510_DIR) + logger.info( + "Activated transformers %s from %s", + TRANSFORMERS_510_VERSION, + _VENV_T5_510_DIR, + ) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_510_DIR + (os.pathsep + _pp if _pp else "") + elif tier == "550": if not _ensure_venv_t5_550_exists(): raise RuntimeError( - f"Cannot activate transformers 5.5.0: " + f"Cannot activate transformers {TRANSFORMERS_550_VERSION}: " f".venv_t5_550 missing at {_VENV_T5_550_DIR}" ) if _VENV_T5_550_DIR not in sys.path: sys.path.insert(0, _VENV_T5_550_DIR) - logger.info("Activated transformers 5.5.0 from %s", _VENV_T5_550_DIR) + logger.info( + "Activated transformers %s from %s", + TRANSFORMERS_550_VERSION, + _VENV_T5_550_DIR, + ) _pp = os.environ.get("PYTHONPATH", "") os.environ["PYTHONPATH"] = _VENV_T5_550_DIR + (os.pathsep + _pp if _pp else "") elif tier == "530": @@ -155,11 +195,10 @@ def activate_transformers_for_subprocess(model_name: str) -> None: def _resolve_base_model(model_name: str) -> str: """If *model_name* points to a LoRA adapter, return its base model. - Checks for ``adapter_config.json`` locally first. Only calls the heavier - ``get_base_model_from_lora`` for paths that are actual local directories - (avoids noisy warnings for plain HF model IDs). - - Returns the original *model_name* unchanged if it is not a LoRA adapter. + Checks ``adapter_config.json`` locally first. Only calls the heavier + ``get_base_model_from_lora`` for real local directories (avoids noisy + warnings for plain HF model IDs). Returns *model_name* unchanged if not a + LoRA adapter. """ # --- Fast local check --------------------------------------------------- local_path = Path(model_name) @@ -201,7 +240,6 @@ def _resolve_base_model(model_name: str) -> str: if local_path.is_dir(): try: from utils.models import get_base_model_from_lora - base = get_base_model_from_lora(model_name) if base: logger.info( @@ -222,11 +260,11 @@ def _resolve_base_model(model_name: str) -> str: def _check_tokenizer_config_needs_v5(model_name: str) -> bool: - """Fetch tokenizer_config.json from HuggingFace and check if the - tokenizer_class requires transformers 5.x. + """True if the model's tokenizer_class requires transformers 5.x. - Results are cached in ``_tokenizer_class_cache`` to avoid repeated fetches. - Returns False on any network/parse error (fail-open to default version). + Checks local tokenizer_config.json, else fetches from HuggingFace. Cached in + ``_tokenizer_class_cache``. Returns False on any network/parse error + (fail-open to default version). """ if model_name in _tokenizer_class_cache: return _tokenizer_class_cache[model_name] @@ -275,59 +313,32 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool: _tokenizer_class_cache[model_name] = result return result except Exception as exc: - logger.debug( - "Could not fetch tokenizer_config.json for '%s': %s", model_name, exc - ) + logger.debug("Could not fetch tokenizer_config.json for '%s': %s", model_name, exc) _tokenizer_class_cache[model_name] = False return False -def _check_config_needs_550(model_name: str) -> bool: - """Check ``config.json`` for architectures or model_type that require - transformers 5.5.0 (e.g. Gemma 4). +def _load_config_json(model_name: str) -> dict | None: + """Return parsed ``config.json`` for *model_name*, checking local files first.""" + if model_name in _config_json_cache: + return _config_json_cache[model_name] - Checks locally first, then falls back to fetching from HuggingFace. - Results are cached in ``_config_needs_550_cache``. - Returns False on any error (fail-open to lower tier). - """ - if model_name in _config_needs_550_cache: - return _config_needs_550_cache[model_name] - - def _check_cfg(cfg: dict) -> bool: - archs = cfg.get("architectures", []) - if any(a in _TRANSFORMERS_550_ARCHITECTURES for a in archs): - return True - if cfg.get("model_type") in _TRANSFORMERS_550_MODEL_TYPES: - return True - return False - - # --- Check local config.json first ------------------------------------ - local_path = Path(model_name) - local_cfg = local_path / "config.json" + local_cfg = Path(model_name) / "config.json" if local_cfg.is_file(): try: with open(local_cfg) as f: cfg = json.load(f) - result = _check_cfg(cfg) - if result: - logger.info( - "Local config.json check: %s needs transformers 5.5.0 " - "(architectures=%s, model_type=%s)", - model_name, - cfg.get("architectures", []), - cfg.get("model_type"), - ) - _config_needs_550_cache[model_name] = result - return result + _config_json_cache[model_name] = cfg + return cfg except Exception as exc: logger.debug("Could not read %s: %s", local_cfg, exc) + _config_json_cache[model_name] = None + return None - # Offline: skip the 10s urllib fetch (fail-open to lower tier). if _env_offline(): - _config_needs_550_cache[model_name] = False - return False + _config_json_cache[model_name] = None + return None - # --- Fall back to fetching from HuggingFace --------------------------- import urllib.request url = f"https://huggingface.co/{model_name}/raw/main/config.json" @@ -335,41 +346,131 @@ def _check_config_needs_550(model_name: str) -> bool: req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) with urllib.request.urlopen(req, timeout = 10) as resp: cfg = json.loads(resp.read().decode()) - result = _check_cfg(cfg) - if result: - logger.info( - "Dynamic config.json check: %s needs transformers 5.5.0 " - "(architectures=%s, model_type=%s)", - model_name, - cfg.get("architectures", []), - cfg.get("model_type"), - ) - _config_needs_550_cache[model_name] = result - return result + _config_json_cache[model_name] = cfg + return cfg except Exception as exc: logger.debug("Could not fetch config.json for '%s': %s", model_name, exc) + _config_json_cache[model_name] = None + return None + + +def _config_matches_tier(cfg: dict, architectures: set[str], model_types: set[str]) -> bool: + archs = cfg.get("architectures", []) + if any(a in architectures for a in archs): + return True + if cfg.get("model_type") in model_types: + return True + return False + + +def _config_needs_550(cfg: dict) -> bool: + return _config_matches_tier( + cfg, + _TRANSFORMERS_550_ARCHITECTURES, + _TRANSFORMERS_550_MODEL_TYPES, + ) + + +def _config_needs_510(cfg: dict) -> bool: + return _config_matches_tier( + cfg, + _TRANSFORMERS_510_ARCHITECTURES, + _TRANSFORMERS_510_MODEL_TYPES, + ) + + +def _check_config_needs_550(model_name: str) -> bool: + """True if ``config.json`` has architectures/model_type needing transformers + 5.5.0 (e.g. Gemma 4). + + Checks locally first, else fetches from HuggingFace. Cached in + ``_config_needs_550_cache``. Returns False on any error (fail-open to lower tier). + """ + if model_name in _config_needs_550_cache: + return _config_needs_550_cache[model_name] + + cfg = _load_config_json(model_name) + if cfg is None: _config_needs_550_cache[model_name] = False return False + result = _config_needs_550(cfg) + if result: + logger.info( + "config.json check: %s needs transformers %s (architectures=%s, model_type=%s)", + model_name, + TRANSFORMERS_550_VERSION, + cfg.get("architectures", []), + cfg.get("model_type"), + ) + _config_needs_550_cache[model_name] = result + return result + + +def _check_config_needs_510(model_name: str) -> bool: + """Check ``config.json`` for Gemma 4 Unified / 12B architectures.""" + if model_name in _config_needs_510_cache: + return _config_needs_510_cache[model_name] + + cfg = _load_config_json(model_name) + if cfg is None: + _config_needs_510_cache[model_name] = False + return False + + result = _config_needs_510(cfg) + if result: + logger.info( + "config.json check: %s needs transformers %s (architectures=%s, model_type=%s)", + model_name, + TRANSFORMERS_510_VERSION, + cfg.get("architectures", []), + cfg.get("model_type"), + ) + _config_needs_510_cache[model_name] = result + return result + def get_transformers_tier(model_name: str) -> str: """Return the transformers tier required for *model_name*. - Returns ``"550"`` for models needing transformers 5.5.0 (e.g. Gemma 4), + Returns ``"510"`` for models needing transformers 5.10.x (Gemma 4 Unified), + ``"550"`` for models needing transformers 5.5.0 (Gemma 4), ``"530"`` for models needing transformers 5.3.0 (e.g. Ministral-3, Qwen3 MoE), or ``"default"`` for everything else (4.57.x). - The 5.5.0 check runs first, then 5.3.0. + Higher 5.x tiers run first. """ lowered = model_name.lower() + # Local checkpoint names can contain architecture substrings in their + # directory names (for example a pytest temp dir). If config.json exists, + # trust it before using name heuristics. + local_cfg = Path(model_name) / "config.json" + if local_cfg.is_file(): + cfg = _load_config_json(model_name) + if cfg is not None and _config_needs_510(cfg): + return "510" + if cfg is not None and _config_needs_550(cfg): + return "550" + if cfg is not None: + local_tc = Path(model_name) / "tokenizer_config.json" + if local_tc.is_file() and _check_tokenizer_config_needs_v5(model_name): + return "530" + return "default" + # --- Fast substring checks (no I/O) ------------------------------------ + if "assistant" in lowered and ("gemma-4" in lowered or "gemma4" in lowered): + return "510" + if any(sub in lowered for sub in TRANSFORMERS_510_MODEL_SUBSTRINGS): + return "510" if any(sub in lowered for sub in TRANSFORMERS_550_MODEL_SUBSTRINGS): return "550" if any(sub in lowered for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS): return "530" - # --- Slow config fallbacks (local file first, then network) ----------- + # --- Slow config fallbacks (network for HF IDs) ------------------------ + if _check_config_needs_510(model_name): + return "510" if _check_config_needs_550(model_name): return "550" if _check_tokenizer_config_needs_v5(model_name): @@ -409,12 +510,11 @@ _PURGE_PREFIXES = ( "trl", "accelerate", "auto_gptq", - # NOTE: bitsandbytes is intentionally EXCLUDED — it registers torch custom - # operators at import time via torch.library.define(). Those registrations - # live in torch's global operator registry which survives module purge. - # Re-importing bitsandbytes after purge → duplicate registration → crash. - # Our own modules that import from transformers at module level - # (e.g. model_config.py: `from transformers import AutoConfig`) + # NOTE: bitsandbytes is intentionally EXCLUDED -- it registers torch custom + # operators via torch.library.define() into torch's global registry, which + # survives module purge; re-importing after purge -> duplicate registration + # -> crash. + # Our own modules that import from transformers at module level. "utils.models", "core.training", "core.inference", @@ -445,6 +545,13 @@ _VENV_T5_530_PACKAGES = ( "tiktoken", ) +_VENV_T5_510_PACKAGES = ( + f"transformers=={TRANSFORMERS_510_VERSION}", + "huggingface_hub==1.8.0", + "hf_xet==1.4.2", + "tiktoken", +) + _VENV_T5_550_PACKAGES = ( f"transformers=={TRANSFORMERS_550_VERSION}", "huggingface_hub==1.8.0", @@ -465,16 +572,15 @@ def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool: pkg_name = parts[0] pkg_version = parts[1] if len(parts) > 1 else None pkg_name_norm = pkg_name.replace("-", "_") - # Check directory exists + # Directory must exist. if not any( - (Path(venv_dir) / d).is_dir() - for d in (pkg_name_norm, pkg_name_norm.replace("_", "-")) + (Path(venv_dir) / d).is_dir() for d in (pkg_name_norm, pkg_name_norm.replace("_", "-")) ): return False - # For unpinned packages, existence is enough + # Unpinned packages: existence is enough. if pkg_version is None: continue - # Check version via .dist-info metadata + # Check version via .dist-info metadata. dist_info_found = False for di in Path(venv_dir).glob(f"{pkg_name_norm}-*.dist-info"): metadata = di / "METADATA" @@ -502,13 +608,13 @@ def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool: def _venv_t5_is_valid() -> bool: - """Backwards-compat: check the 5.5.0 venv.""" + """Backwards-compat: check the Gemma 4 sidecar venv.""" return _venv_dir_is_valid(_VENV_T5_550_DIR, _VENV_T5_550_PACKAGES) def _install_to_dir(pkg: str, target_dir: str) -> bool: """Install a single package into *target_dir*, preferring uv then pip.""" - # Try uv first (faster) if already on PATH -- do NOT install uv at runtime + # Try uv first (faster) if on PATH -- do NOT install uv at runtime. if shutil.which("uv"): result = subprocess.run( [ @@ -533,7 +639,7 @@ def _install_to_dir(pkg: str, target_dir: str) -> bool: return True logger.warning("uv install of %s failed, falling back to pip", pkg) - # Fallback to pip + # Fallback to pip. result = subprocess.run( [ sys.executable, @@ -563,9 +669,7 @@ def _ensure_venv_dir(venv_dir: str, packages: tuple[str, ...], label: str) -> bo if _venv_dir_is_valid(venv_dir, packages): return True - logger.warning( - "%s not found or incomplete at %s -- installing at runtime", label, venv_dir - ) + logger.warning("%s not found or incomplete at %s -- installing at runtime", label, venv_dir) shutil.rmtree(venv_dir, ignore_errors = True) os.makedirs(venv_dir, exist_ok = True) for pkg in packages: @@ -577,20 +681,29 @@ def _ensure_venv_dir(venv_dir: str, packages: tuple[str, ...], label: str) -> bo def _ensure_venv_t5_530_exists() -> bool: """Ensure .venv_t5_530/ exists with transformers 5.3.0.""" - return _ensure_venv_dir( - _VENV_T5_530_DIR, _VENV_T5_530_PACKAGES, "transformers 5.3.0" - ) + return _ensure_venv_dir(_VENV_T5_530_DIR, _VENV_T5_530_PACKAGES, "transformers 5.3.0") def _ensure_venv_t5_550_exists() -> bool: """Ensure .venv_t5_550/ exists with transformers 5.5.0.""" return _ensure_venv_dir( - _VENV_T5_550_DIR, _VENV_T5_550_PACKAGES, "transformers 5.5.0" + _VENV_T5_550_DIR, + _VENV_T5_550_PACKAGES, + f"transformers {TRANSFORMERS_550_VERSION}", + ) + + +def _ensure_venv_t5_510_exists() -> bool: + """Ensure .venv_t5_510/ exists with transformers 5.10.x.""" + return _ensure_venv_dir( + _VENV_T5_510_DIR, + _VENV_T5_510_PACKAGES, + f"transformers {TRANSFORMERS_510_VERSION}", ) def _ensure_venv_t5_exists() -> bool: - """Backwards-compat: ensure the 5.5.0 venv exists.""" + """Backwards-compat: ensure the Gemma 4 5.5 sidecar venv exists.""" return _ensure_venv_t5_550_exists() @@ -610,7 +723,7 @@ def _activate_venv(venv_dir: str, label: str) -> None: def _deactivate_5x() -> None: """Remove all .venv_t5_*/ dirs from sys.path, purge stale modules, reimport.""" - for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR): + for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR): while d in sys.path: sys.path.remove(d) logger.info("Removed venv_t5 dirs from sys.path") @@ -626,22 +739,28 @@ def _deactivate_5x() -> None: def ensure_transformers_version(model_name: str) -> None: """Ensure the correct ``transformers`` version is active for *model_name*. - Uses sys.path with .venv_t5_530/ or .venv_t5_550/ (pre-installed by setup.sh): + Uses sys.path with .venv_t5_510/, .venv_t5_550/, or .venv_t5_530/ + (pre-installed by setup.sh): + • Need 5.10.x → prepend .venv_t5_510/ to sys.path, purge modules. • Need 5.5.0 → prepend .venv_t5_550/ to sys.path, purge modules. • Need 5.3.0 → prepend .venv_t5_530/ to sys.path, purge modules. • Need 4.x → remove all .venv_t5_*/ from sys.path, purge modules. - For LoRA adapters with custom names, the base model is resolved from + For custom-named LoRA adapters, the base model is resolved from ``adapter_config.json`` before checking. - NOTE: Training and inference use subprocess isolation instead of this - function. This is only used by the export path (routes/export.py). + NOTE: Training and inference use subprocess isolation instead. Used only by + the export path (routes/export.py). """ - # Resolve LoRA adapters to their base model for accurate detection + # Resolve LoRA adapters to their base model for accurate detection. resolved = _resolve_base_model(model_name) tier = get_transformers_tier(resolved) - if tier == "550": + if tier == "510": + target_version = TRANSFORMERS_510_VERSION + venv_dir = _VENV_T5_510_DIR + ensure_fn = _ensure_venv_t5_510_exists + elif tier == "550": target_version = TRANSFORMERS_550_VERSION venv_dir = _VENV_T5_550_DIR ensure_fn = _ensure_venv_t5_550_exists @@ -676,10 +795,10 @@ def ensure_transformers_version(model_name: str) -> None: model_name, ) return - # Different 5.x → need to switch (e.g. 5.3.0 loaded but need 5.5.0) + # Different 5.x -> need to switch (e.g. 5.3.0 loaded but need 5.10.x). in_memory_major = int(in_memory.split(".")[0]) if in_memory_major == target_major and venv_dir is None: - # Both are default (4.x) — close enough + # Both are default (4.x) — close enough. logger.info( "transformers %s already loaded — correct for '%s'", in_memory, @@ -689,19 +808,16 @@ def ensure_transformers_version(model_name: str) -> None: # --- Switch version ----------------------------------------------------- if venv_dir is not None: - # First remove any other 5.x venv from sys.path + # First remove any other 5.x venv from sys.path. _deactivate_5x() if not ensure_fn(): raise RuntimeError( - f"Cannot activate transformers {target_version}: " - f"venv missing at {venv_dir}" + f"Cannot activate transformers {target_version}: " f"venv missing at {venv_dir}" ) logger.info("Activating transformers %s…", target_version) _activate_venv(venv_dir, f"transformers {target_version}") else: - logger.info( - "Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION - ) + logger.info("Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION) _deactivate_5x() final = _get_in_memory_version() diff --git a/studio/backend/utils/update_status.py b/studio/backend/utils/update_status.py index 9142203a69..ad9dabcf36 100644 --- a/studio/backend/utils/update_status.py +++ b/studio/backend/utils/update_status.py @@ -3,9 +3,8 @@ """Web update status helpers for browser-served Unsloth Studio. -This module is intentionally side-effect light: no network work happens at -import time or from /api/health. The PyPI check is lazy, cached, and only used -for normal PyPI-managed installs. +Side-effect light: no network work at import time or from /api/health. +The PyPI check is lazy, cached, and only for PyPI-managed installs. """ from __future__ import annotations @@ -66,18 +65,14 @@ def reset_update_status_cache() -> None: def detect_install_source() -> str: """Return a coarse install source without exposing local paths. - Sources are intentionally conservative. PEP 610 local/vcs metadata wins. - Legacy source installs are treated as local only when package files resolve - outside site-packages/dist-packages and under a Git checkout. + Conservative: PEP 610 local/vcs metadata wins. Legacy source + installs count as local only when package files resolve outside + site-packages/dist-packages and under a Git checkout. """ try: dist = distribution(PACKAGE_NAME) except PackageNotFoundError: - return ( - "local_repo" - if _path_has_git_parent(_repo_root_from_this_file()) - else "unknown" - ) + return "local_repo" if _path_has_git_parent(_repo_root_from_this_file()) else "unknown" try: direct_url = dist.read_text("direct_url.json") @@ -146,9 +141,7 @@ def get_studio_update_status(current_version: str) -> dict[str, Any]: current_version = current_version, latest_version = None, install_source = install_source, - reason = "invalid_current_version" - if current_version != "dev" - else "dev_build", + reason = "invalid_current_version" if current_version != "dev" else "dev_build", ) latest_result = get_latest_pypi_version() if latest_result.latest_version is None: @@ -216,9 +209,7 @@ def get_latest_pypi_version() -> LatestVersionResult: error = "Could not check PyPI update metadata.", ) - ttl = ( - PYPI_SUCCESS_TTL_SECONDS if result.latest_version else PYPI_FAILURE_TTL_SECONDS - ) + ttl = PYPI_SUCCESS_TTL_SECONDS if result.latest_version else PYPI_FAILURE_TTL_SECONDS with _cache_condition: _latest_version_cache = _LatestVersionCacheEntry( result = result, @@ -262,9 +253,7 @@ def _fetch_latest_pypi_version() -> LatestVersionResult: error = "Could not reach PyPI for update metadata.", ) - latest = ( - payload.get("info", {}).get("version") if isinstance(payload, dict) else None - ) + latest = payload.get("info", {}).get("version") if isinstance(payload, dict) else None if not isinstance(latest, str) or not latest.strip(): return LatestVersionResult( latest_version = None, @@ -366,9 +355,4 @@ def _parse_current_version(current_version: str) -> Version | None: def _utc_now_iso() -> str: - return ( - datetime.now(timezone.utc) - .replace(microsecond = 0) - .isoformat() - .replace("+00:00", "Z") - ) + return datetime.now(timezone.utc).replace(microsecond = 0).isoformat().replace("+00:00", "Z") diff --git a/studio/backend/utils/upload_limits.py b/studio/backend/utils/upload_limits.py new file mode 100644 index 0000000000..c21ea69af7 --- /dev/null +++ b/studio/backend/utils/upload_limits.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared Studio upload/request size limits.""" + +from __future__ import annotations + +import os +from typing import Any + +UPLOAD_LIMIT_SETTING_KEY = "max_upload_size_mb" +DEFAULT_UPLOAD_LIMIT_MB = 500 +MIN_UPLOAD_LIMIT_MB = 1 +MAX_UPLOAD_LIMIT_MB = 8192 +_BYTES_PER_MB = 1024 * 1024 +MULTIPART_OVERHEAD_BYTES = 10 * _BYTES_PER_MB + +LOCAL_SEED_UPLOAD_MAX_BYTES = 100 * _BYTES_PER_MB +LOCAL_SEED_UPLOAD_MAX_LABEL = "100MB" +UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES = 500 * _BYTES_PER_MB +UNSTRUCTURED_RECIPE_UPLOAD_MAX_LABEL = "500MB" +UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES = 1024 * _BYTES_PER_MB +UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_LABEL = "1GB" + + +def _coerce_upload_limit_mb(value: Any) -> int | None: + if isinstance(value, bool): + return None + try: + parsed = int(value) + except (TypeError, ValueError): + return None + if parsed < MIN_UPLOAD_LIMIT_MB or parsed > MAX_UPLOAD_LIMIT_MB: + return None + return parsed + + +def default_upload_limit_mb() -> int: + env_value = _coerce_upload_limit_mb(os.environ.get("UNSLOTH_STUDIO_MAX_BODY_MB")) + return env_value or DEFAULT_UPLOAD_LIMIT_MB + + +def validate_upload_limit_mb(value: Any) -> int: + parsed = _coerce_upload_limit_mb(value) + if parsed is None: + raise ValueError( + f"Upload limit must be a whole number from {MIN_UPLOAD_LIMIT_MB} to {MAX_UPLOAD_LIMIT_MB} MB." + ) + return parsed + + +def get_upload_limit_mb() -> int: + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(UPLOAD_LIMIT_SETTING_KEY, None) + except Exception: + stored = None + return _coerce_upload_limit_mb(stored) or default_upload_limit_mb() + + +def set_upload_limit_mb(value: Any) -> int: + parsed = validate_upload_limit_mb(value) + from storage.studio_db import upsert_app_settings + + upsert_app_settings({UPLOAD_LIMIT_SETTING_KEY: parsed}) + return parsed + + +def upload_limit_bytes(limit_mb: int | None = None) -> int: + return (limit_mb if limit_mb is not None else get_upload_limit_mb()) * _BYTES_PER_MB + + +def get_upload_limit_bytes() -> int: + return upload_limit_bytes() + + +def upload_limit_label(limit_mb: int | None = None) -> str: + return f"{limit_mb if limit_mb is not None else get_upload_limit_mb()}MB" + + +def get_upload_limit_label() -> str: + return upload_limit_label() + + +def default_request_body_limit_bytes() -> int: + """Default protected-route body cap for non-upload requests.""" + + return default_upload_limit_mb() * _BYTES_PER_MB + + +def upload_request_limit_bytes(file_limit_bytes: int | None = None) -> int: + """Request cap for upload routes, including multipart field overhead.""" + + return ( + file_limit_bytes if file_limit_bytes is not None else get_upload_limit_bytes() + ) + MULTIPART_OVERHEAD_BYTES diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 4e61a5b969..3818253ac9 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Shared backend utilities -""" +"""Shared backend utilities.""" import os import structlog @@ -17,17 +15,69 @@ import tempfile logger = get_logger(__name__) +# ── Client-safe error helpers ─────────────────────────────────── +# Never return raw exception text to clients; log server-side, return generic. + + +def safe_error_detail(error: Exception, fallback: str = "An internal error occurred") -> str: + """Map an exception to a generic, client-safe message (never raw + ``str(error)``, which can leak paths). Log the real exception server-side. + """ + text = str(error).lower() + if ( + isinstance(error, (ConnectionError, TimeoutError)) + or "connection" in text + or "timed out" in text + or "timeout" in text + ): + return "Could not reach an upstream service. Please try again." + if "out of memory" in text or "cuda error" in text: + return "Ran out of memory. Try a smaller model or shorter input." + return fallback + + +def safe_curated_detail(error: Exception, fallback: str = "An internal error occurred") -> str: + """Client-safe text for curated domain/validation exceptions. + + Keeps the message (paths stripped) instead of a generic fallback; for known + exception types only (use ``safe_error_detail`` for generic ``Exception``). + """ + from utils.native_path_leases import redact_native_paths + + msg = redact_native_paths(str(error)).strip() + return msg or fallback + + +def log_and_http_error( + error: Exception, + status_code: int, + public_message: str, + *, + event: str = "request_failed", + log = None, +): + """Log ``error`` in full server-side and return an ``HTTPException`` whose + ``detail`` is only ``public_message`` -- never the raw exception text. + + Usage: raise log_and_http_error(e, 500, "Failed to start training") + """ + from fastapi import HTTPException + + # exc_info=error works for both structlog and stdlib loggers. + (log or logger).error(f"{event}: {error}", exc_info = error) + return HTTPException(status_code = status_code, detail = public_message) + + @contextmanager def without_hf_auth(): """ - Context manager to temporarily disable HuggingFace authentication. + Temporarily disable HuggingFace authentication. Usage: with without_hf_auth(): # Code that should run without cached tokens model_info(model_name, token=None) """ - # Save environment variables saved_env = {} env_vars = ["HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_HOME"] for var in env_vars: @@ -35,11 +85,10 @@ def without_hf_auth(): saved_env[var] = os.environ[var] del os.environ[var] - # Save disable flag saved_disable = os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN") os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "1" - # Move token files temporarily + # Move token files aside temporarily token_files = [] token_locations = [ Path.home() / ".cache" / "huggingface" / "token", @@ -64,7 +113,7 @@ def without_hf_auth(): except Exception as e: logger.error(f"Failed to restore token {original}: {e}") - # Restore environment + # Restore env for var, value in saved_env.items(): os.environ[var] = value @@ -76,14 +125,11 @@ def without_hf_auth(): def format_error_message(error: Exception, model_name: str) -> str: """ - Format user-friendly error messages for common issues. + Format a user-friendly error message for common load issues. Args: error: The exception that occurred model_name: Name of the model being loaded - - Returns: - User-friendly error string """ error_str = str(error).lower() model_short = model_name.split("/")[-1] if "/" in model_name else model_name @@ -114,5 +160,4 @@ def format_error_message(error: Exception, model_name: str) -> str: ) return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory." - # Generic fallback return str(error) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 5c42e890d1..98697df83c 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -15,26 +15,21 @@ import urllib.request from typing import Callable from utils.native_path_leases import child_env_without_native_path_secret +from utils.subprocess_compat import windows_hidden_subprocess_kwargs _logger = logging.getLogger(__name__) -FLASH_ATTN_RELEASE_BASE_URL = ( - "https://github.com/Dao-AILab/flash-attention/releases/download" -) +FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/releases/download" @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, ...). + """Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell). - 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. + Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels + fail to load, so callers use this to skip the flash-attn install path. Cached + for the process lifetime; tests mocking nvidia-smi must call + ``has_blackwell_gpu.cache_clear()`` first. """ exe = shutil.which("nvidia-smi") if not exe: @@ -106,6 +101,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non text = True, timeout = timeout, env = child_env_without_native_path_secret(), + **windows_hidden_subprocess_kwargs(), ) except subprocess.TimeoutExpired: return None diff --git a/studio/frontend/.npmrc b/studio/frontend/.npmrc index 8e21abe7a2..f2d15a4f15 100644 --- a/studio/frontend/.npmrc +++ b/studio/frontend/.npmrc @@ -6,11 +6,8 @@ # upstream removal. npm interprets the bare integer as DAYS; do not # append `d`, npm 11.x will parse `7d` as a Date string and abort. min-release-age=7 -# Defensive alias: `minimum-release-age` takes minutes (10080 = 7 days). -# Some npm versions / wrappers consult one key but not the other; setting -# both means a single setting-name parse change upstream cannot silently -# disable the cooldown. The two keys MUST agree; do not let them drift. -minimum-release-age=10080 +# Do not re-add the old `minimum-release-age` alias: npm >=11.16 warns on +# unknown project configs and npm 12 stops accepting them. # Belt-and-braces: refuse to write back loose `^x.y.z` ranges into # package.json when a maintainer runs `npm install ` locally. This # does NOT rewrite already-present ranges (those need an explicit diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json index 80f5d0a701..a521552d79 100644 --- a/studio/frontend/package-lock.json +++ b/studio/frontend/package-lock.json @@ -31,12 +31,14 @@ "@tailwindcss/vite": "^4.2.2", "@tanstack/react-router": "1.169.2", "@tanstack/react-table": "^8.21.3", + "@tanstack/react-virtual": "3.13.25", "@tauri-apps/api": "^2.10.1", "@tauri-apps/plugin-clipboard-manager": "^2.3.2", "@tauri-apps/plugin-notification": "^2.3.3", "@tauri-apps/plugin-opener": "^2.5.3", "@tauri-apps/plugin-process": "^2.3.1", "@tauri-apps/plugin-updater": "^2.10.1", + "@tauri-apps/plugin-window-state": "^2.4.1", "@toolwind/corner-shape": "^0.0.8-3", "@xyflow/react": "^12.10.0", "assistant-stream": "0.3.12", @@ -57,6 +59,7 @@ "react": "^19.2.4", "react-day-picker": "^9.13.2", "react-dom": "^19.2.4", + "react-pdf": "10.4.1", "react-resizable-panels": "^4.6.4", "recharts": "3.7.0", "shadcn": "^4.2.0", @@ -855,41 +858,10 @@ "integrity": "sha512-jigsZK+sMF/cuiB7sERuo9V7N9jx+dhmHHnQyDSVdpZwVutaBu7WvNYqMDLSgFgfB30n452TP3vjDAvFC973mA==", "license": "MIT" }, - "node_modules/@chevrotain/cst-dts-gen": { - "version": "12.0.0", - "resolved": "https://registry.npmjs.org/@chevrotain/cst-dts-gen/-/cst-dts-gen-12.0.0.tgz", - "integrity": "sha512-fSL4KXjTl7cDgf0B5Rip9Q05BOrYvkJV/RrBTE/bKDN096E4hN/ySpcBK5B24T76dlQ2i32Zc3PAE27jFnFrKg==", - "license": "Apache-2.0", - "dependencies": { - "@chevrotain/gast": "12.0.0", - "@chevrotain/types": "12.0.0" - } - }, - "node_modules/@chevrotain/gast": { - "version": "12.0.0", - "resolved": "https://registry.npmjs.org/@chevrotain/gast/-/gast-12.0.0.tgz", - "integrity": "sha512-1ne/m3XsIT8aEdrvT33so0GUC+wkctpUPK6zU9IlOyJLUbR0rg4G7ZiApiJbggpgPir9ERy3FRjT6T7lpgetnQ==", - "license": "Apache-2.0", - "dependencies": { - "@chevrotain/types": "12.0.0" - } - }, - "node_modules/@chevrotain/regexp-to-ast": { - "version": "12.0.0", - "resolved": "https://registry.npmjs.org/@chevrotain/regexp-to-ast/-/regexp-to-ast-12.0.0.tgz", - "integrity": "sha512-p+EW9MaJwgaHguhoqwOtx/FwuGr+DnNn857sXWOi/mClXIkPGl3rn7hGNWvo31HA3vyeQxjqe+H36yZJwYU8cA==", - "license": "Apache-2.0" - }, "node_modules/@chevrotain/types": { - "version": "12.0.0", - "resolved": "https://registry.npmjs.org/@chevrotain/types/-/types-12.0.0.tgz", - "integrity": "sha512-S+04vjFQKeuYw0/eW3U52LkAHQsB1ASxsPGsLPUyQgrZ2iNNibQrsidruDzjEX2JYfespXMG0eZmXlhA6z7nWA==", - "license": "Apache-2.0" - }, - "node_modules/@chevrotain/utils": { - "version": "12.0.0", - "resolved": "https://registry.npmjs.org/@chevrotain/utils/-/utils-12.0.0.tgz", - "integrity": "sha512-lB59uJoaGIfOOL9knQqQRfhl9g7x8/wqFkp13zTdkRu1huG9kg6IJs1O8hqj9rs6h7orGxHJUKb+mX3rPbWGhA==", + "version": "11.1.2", + "resolved": "https://registry.npmjs.org/@chevrotain/types/-/types-11.1.2.tgz", + "integrity": "sha512-U+HFai5+zmJCkK86QsaJtoITlboZHBqrVketcO2ROv865xfCMSFpELQoz1GkX5GzME8pTa+3kbKrZHQtI0gdbw==", "license": "Apache-2.0" }, "node_modules/@dagrejs/dagre": { @@ -1597,12 +1569,12 @@ } }, "node_modules/@mermaid-js/parser": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@mermaid-js/parser/-/parser-1.1.0.tgz", - "integrity": "sha512-gxK9ZX2+Fex5zu8LhRQoMeMPEHbc73UKZ0FQ54YrQtUxE1VVhMwzeNtKRPAu5aXks4FasbMe4xB4bWrmq6Jlxw==", + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@mermaid-js/parser/-/parser-1.1.1.tgz", + "integrity": "sha512-VuHdsYMK1bT6X2JbcAaWAhugTRvRBRyuZgd+c22swUeI9g/ntaxF7CY7dYarhZovofCbUNO0G7JesfmNtjYOCw==", "license": "MIT", "dependencies": { - "langium": "^4.0.0" + "@chevrotain/types": "~11.1.1" } }, "node_modules/@modelcontextprotocol/sdk": { @@ -1690,6 +1662,256 @@ "integrity": "sha512-CecwLWx3rhxVQF6V4bAgPS5t+So2sTbPgAzafKkVizyi7tlwpcFpdFqq+wqF2OwNBmqFuu6tOyouTuxgpMfzmA==", "license": "MIT" }, + "node_modules/@napi-rs/canvas": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas/-/canvas-0.1.100.tgz", + "integrity": "sha512-xglYA6q3XO5P3BNJYxVZ1IV7DLVjp1Py6nwag88YntrS+3vKHyYcMqXVS4ZztJmwz2uGvz1FWhI/4LgbR5uQDA==", + "license": "MIT", + "optional": true, + "workspaces": [ + "e2e/*" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + }, + "optionalDependencies": { + "@napi-rs/canvas-android-arm64": "0.1.100", + "@napi-rs/canvas-darwin-arm64": "0.1.100", + "@napi-rs/canvas-darwin-x64": "0.1.100", + "@napi-rs/canvas-linux-arm-gnueabihf": "0.1.100", + "@napi-rs/canvas-linux-arm64-gnu": "0.1.100", + "@napi-rs/canvas-linux-arm64-musl": "0.1.100", + "@napi-rs/canvas-linux-riscv64-gnu": "0.1.100", + "@napi-rs/canvas-linux-x64-gnu": "0.1.100", + "@napi-rs/canvas-linux-x64-musl": "0.1.100", + "@napi-rs/canvas-win32-arm64-msvc": "0.1.100", + "@napi-rs/canvas-win32-x64-msvc": "0.1.100" + } + }, + "node_modules/@napi-rs/canvas-android-arm64": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-android-arm64/-/canvas-android-arm64-0.1.100.tgz", + "integrity": "sha512-hjhCKhntPv9+t4ckHymdx0phYNcVW+GKQR6Lzw2zE+pOVjOplSmtx9nNNknTjbEDLcuLZqA1y8ufKg1XfgftzQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-darwin-arm64": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-arm64/-/canvas-darwin-arm64-0.1.100.tgz", + "integrity": "sha512-2PcswRaC7Ly645DGt88///zuFDhJxJYdKAs1uU3mfk1atYkXufgcgLfBpk6Tm12nCQBaNt1wpybuPZ4qOhTo8A==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-darwin-x64": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-x64/-/canvas-darwin-x64-0.1.100.tgz", + "integrity": "sha512-ePNZtj7pNIva/siZMg+HmbeozkIjqUIYdoymH8HaA3qK7LfzFN4WMBM8G6HQ9ZC+H3+Dnn5pqtiXpgLykaPOhw==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-linux-arm-gnueabihf": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm-gnueabihf/-/canvas-linux-arm-gnueabihf-0.1.100.tgz", + "integrity": "sha512-d5cDB48oWFGU8/XPhUOFAlySgb/VAu7D+s8fi55K1Pcfg8aPplHWqMgibhVLU8ky7Pyg/fuiVLz4Nf3JrSTuUA==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-linux-arm64-gnu": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-gnu/-/canvas-linux-arm64-gnu-0.1.100.tgz", + "integrity": "sha512-rDxgxRu69RvDlX/bh9o22DxLsGr8EqsNgotL9+RwQE1S0b0cqeatqsw6aW45mukm0B42DIAaAacKaYQ8cqS1nw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-linux-arm64-musl": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-musl/-/canvas-linux-arm64-musl-0.1.100.tgz", + "integrity": "sha512-K3mDW66N+xT2/V439u1alFANiBUjdEx2gLiNYnCmUsva5jZMxWTjafBYwTzYK+EMFMHrUoabuU+T1BIP5CgbYQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-linux-riscv64-gnu": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-riscv64-gnu/-/canvas-linux-riscv64-gnu-0.1.100.tgz", + "integrity": "sha512-mooqUBTIsccZpnoQC4NgrC1v6C1vof39etLNMnBwCY+p0gajWJvAHLGQ6g/gGyS5YrpDW+GefSN4+Cvcr08UWw==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-linux-x64-gnu": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-gnu/-/canvas-linux-x64-gnu-0.1.100.tgz", + "integrity": "sha512-1eCvkDCazm7FFhsT7DfGOdSaHgZVK3bt/dSBl5EWHOWmnz+I7j8tPseJqqD81NF+MH21jKUK4wQSDjN0mdhnTg==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-linux-x64-musl": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-musl/-/canvas-linux-x64-musl-0.1.100.tgz", + "integrity": "sha512-20arT6lnI19S68qNlii73TSEDbECNgzMz2EpldC1V3mZFuRkeujXkcebRk0LRJe9SEUAooYiLokfMViY8IX7yA==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-win32-arm64-msvc": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-arm64-msvc/-/canvas-win32-arm64-msvc-0.1.100.tgz", + "integrity": "sha512-DZFFT1wIAg37LJw37yhMRFfjATd3vTQzjZ1Yki8u2vhO6Hi5VE6BVaGQ1aaDu7xb4iMErz+9EOwjpS7xcxFeBw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-win32-x64-msvc": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-x64-msvc/-/canvas-win32-x64-msvc-0.1.100.tgz", + "integrity": "sha512-MyT1j3mHC2+Lu4pBi9mKyMJhtP6U7k7EldY7sj/uS5gJA65gTXt8MefJQXLJo5d/vZbuWmfxzkEUNc/urV3pHA==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, "node_modules/@napi-rs/wasm-runtime": { "version": "1.1.4", "resolved": "https://registry.npmjs.org/@napi-rs/wasm-runtime/-/wasm-runtime-1.1.4.tgz", @@ -6142,6 +6364,23 @@ "react-dom": ">=16.8" } }, + "node_modules/@tanstack/react-virtual": { + "version": "3.13.25", + "resolved": "https://registry.npmjs.org/@tanstack/react-virtual/-/react-virtual-3.13.25.tgz", + "integrity": "sha512-bmNoqMu6gcAW9JGrKVB0Q1tN1i5RONZF8r1fW0bbE4Oyf3DwEGnzzQJ2OW+Ozg1P4s8PyugkHg2ULZoFQN+cqw==", + "license": "MIT", + "dependencies": { + "@tanstack/virtual-core": "3.15.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + }, + "peerDependencies": { + "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0", + "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" + } + }, "node_modules/@tanstack/router-core": { "version": "1.169.2", "resolved": "https://registry.npmjs.org/@tanstack/router-core/-/router-core-1.169.2.tgz", @@ -6184,6 +6423,16 @@ "url": "https://github.com/sponsors/tannerlinsley" } }, + "node_modules/@tanstack/virtual-core": { + "version": "3.15.0", + "resolved": "https://registry.npmjs.org/@tanstack/virtual-core/-/virtual-core-3.15.0.tgz", + "integrity": "sha512-0AwPGx0I8QxPYjAxShT/+z+ZOe9u8mW5rsXvivCTjRfRmz9a43+3mRyi4wwlyoUqOC56q/jatKa0Bh9M99BEHQ==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + } + }, "node_modules/@tauri-apps/api": { "version": "2.11.0", "resolved": "https://registry.npmjs.org/@tauri-apps/api/-/api-2.11.0.tgz", @@ -6239,6 +6488,15 @@ "@tauri-apps/api": "^2.10.1" } }, + "node_modules/@tauri-apps/plugin-window-state": { + "version": "2.4.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/plugin-window-state/-/plugin-window-state-2.4.1.tgz", + "integrity": "sha512-OuvdrzyY8Q5Dbzpj+GcrnV1iCeoZbcFdzMjanZMMcAEUNy/6PH5pxZPXpaZLOR7whlzXiuzx0L9EKZbH7zpdRw==", + "license": "MIT OR Apache-2.0", + "dependencies": { + "@tauri-apps/api": "^2.8.0" + } + }, "node_modules/@toolwind/corner-shape": { "version": "0.0.8-3", "resolved": "https://registry.npmjs.org/@toolwind/corner-shape/-/corner-shape-0.0.8-3.tgz", @@ -6284,9 +6542,9 @@ } }, "node_modules/@ts-morph/common/node_modules/brace-expansion": { - "version": "5.0.5", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.5.tgz", - "integrity": "sha512-VZznLgtwhn+Mact9tfiwx64fA9erHH/MCXEUfB/0bX/6Fz6ny5EGTXYltMocqg4xFAQZtnO3DHWWXi8RiuN7cQ==", + "version": "5.0.6", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", + "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", "license": "MIT", "dependencies": { "balanced-match": "^4.0.2" @@ -6932,9 +7190,9 @@ } }, "node_modules/@typescript-eslint/typescript-estree/node_modules/brace-expansion": { - "version": "5.0.5", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.5.tgz", - "integrity": "sha512-VZznLgtwhn+Mact9tfiwx64fA9erHH/MCXEUfB/0bX/6Fz6ny5EGTXYltMocqg4xFAQZtnO3DHWWXi8RiuN7cQ==", + "version": "5.0.6", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", + "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", "dev": true, "license": "MIT", "dependencies": { @@ -7627,34 +7885,6 @@ "url": "https://github.com/sponsors/wooorm" } }, - "node_modules/chevrotain": { - "version": "12.0.0", - "resolved": "https://registry.npmjs.org/chevrotain/-/chevrotain-12.0.0.tgz", - "integrity": "sha512-csJvb+6kEiQaqo1woTdSAuOWdN0WTLIydkKrBnS+V5gZz0oqBrp4kQ35519QgK6TpBThiG3V1vNSHlIkv4AglQ==", - "license": "Apache-2.0", - "dependencies": { - "@chevrotain/cst-dts-gen": "12.0.0", - "@chevrotain/gast": "12.0.0", - "@chevrotain/regexp-to-ast": "12.0.0", - "@chevrotain/types": "12.0.0", - "@chevrotain/utils": "12.0.0" - }, - "engines": { - "node": ">=22.0.0" - } - }, - "node_modules/chevrotain-allstar": { - "version": "0.4.3", - "resolved": "https://registry.npmjs.org/chevrotain-allstar/-/chevrotain-allstar-0.4.3.tgz", - "integrity": "sha512-2X4mkroolSMKqW+H22pyPMUVDqYZzPhephTmg/NODKb1IGYPHfxfhcW0EjS7wcPJNbze2i4vBWT7zT5FKF2lrQ==", - "license": "MIT", - "dependencies": { - "lodash-es": "^4.18.1" - }, - "peerDependencies": { - "chevrotain": "^12.0.0" - } - }, "node_modules/class-variance-authority": { "version": "0.7.1", "resolved": "https://registry.npmjs.org/class-variance-authority/-/class-variance-authority-0.7.1.tgz", @@ -10034,9 +10264,9 @@ } }, "node_modules/hono": { - "version": "4.12.17", - "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.17.tgz", - "integrity": "sha512-FbJJNb/XgX7YW0hX/V8w5oYLztKEsRLykCMZWt1WdLtsfjzMvmoqWBA4H4t5norinq8/rh20oiZYr+WSl4UzAQ==", + "version": "4.12.21", + "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.21.tgz", + "integrity": "sha512-uV63apnb0kyPtAUwoWgaGh9HyIFcv8lgmzPZSiTBQAFOFGIzka5EZ1dZocmGnn0XdX0+XTqJ6Tqv7selMuGLRQ==", "license": "MIT", "engines": { "node": ">=16.9.0" @@ -10189,9 +10419,9 @@ } }, "node_modules/ip-address": { - "version": "10.1.0", - "resolved": "https://registry.npmjs.org/ip-address/-/ip-address-10.1.0.tgz", - "integrity": "sha512-XXADHxXmvT9+CRxhXg56LJovE+bmWnEWB78LB83VZTprKTmaC5QfruXocxzTZ2Kl0DNwKuBdlIhjL8LeY8Sf8Q==", + "version": "10.1.1", + "resolved": "https://registry.npmjs.org/ip-address/-/ip-address-10.1.1.tgz", + "integrity": "sha512-1FMu8/N15Ck1BL551Jf42NYIoin2unWjLQ2Fze/DXryJRl5twqtwNHlO39qERGbIOcKYWHdgRryhOC+NG4eaLw==", "license": "MIT", "engines": { "node": ">= 12" @@ -10617,24 +10847,6 @@ "node": ">=6" } }, - "node_modules/langium": { - "version": "4.2.3", - "resolved": "https://registry.npmjs.org/langium/-/langium-4.2.3.tgz", - "integrity": "sha512-sOPIi4hISFnY7twwV97ca1TsxpBtXq0URu/LL1AvxwccPG/RIBBlKS7a/f/EL6w8lTNaS0EFs/F+IdSOaqYpng==", - "license": "MIT", - "dependencies": { - "@chevrotain/regexp-to-ast": "~12.0.0", - "chevrotain": "~12.0.0", - "chevrotain-allstar": "~0.4.3", - "vscode-languageserver": "~9.0.1", - "vscode-languageserver-textdocument": "~1.0.11", - "vscode-uri": "~3.1.0" - }, - "engines": { - "node": ">=20.10.0", - "npm": ">=10.2.3" - } - }, "node_modules/layout-base": { "version": "1.0.2", "resolved": "https://registry.npmjs.org/layout-base/-/layout-base-1.0.2.tgz", @@ -10998,6 +11210,18 @@ "url": "https://github.com/sponsors/wooorm" } }, + "node_modules/loose-envify": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/loose-envify/-/loose-envify-1.4.0.tgz", + "integrity": "sha512-lyuxPGr/Wfhrlem2CL/UcnUc1zcqKAImBDzukY7Y5F/yQiNdko6+fRLevlw1HgMySw7f611UIY408EtxRSoK3Q==", + "license": "MIT", + "dependencies": { + "js-tokens": "^3.0.0 || ^4.0.0" + }, + "bin": { + "loose-envify": "cli.js" + } + }, "node_modules/lop": { "version": "0.4.2", "resolved": "https://registry.npmjs.org/lop/-/lop-0.4.2.tgz", @@ -11036,6 +11260,24 @@ "@jridgewell/sourcemap-codec": "^1.5.5" } }, + "node_modules/make-cancellable-promise": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/make-cancellable-promise/-/make-cancellable-promise-2.0.0.tgz", + "integrity": "sha512-3SEQqTpV9oqVsIWqAcmDuaNeo7yBO3tqPtqGRcKkEo0lrzD3wqbKG9mkxO65KoOgXqj+zH2phJ2LiAsdzlogSw==", + "license": "MIT", + "funding": { + "url": "https://github.com/wojtekmaj/make-cancellable-promise?sponsor=1" + } + }, + "node_modules/make-event-props": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/make-event-props/-/make-event-props-2.0.0.tgz", + "integrity": "sha512-G/hncXrl4Qt7mauJEXSg3AcdYzmpkIITTNl5I+rH9sog5Yw0kK6vseJjCaPfOXqOqQuPUP89Rkhfz5kPS8ijtw==", + "license": "MIT", + "funding": { + "url": "https://github.com/wojtekmaj/make-event-props?sponsor=1" + } + }, "node_modules/mammoth": { "version": "1.12.0", "resolved": "https://registry.npmjs.org/mammoth/-/mammoth-1.12.0.tgz", @@ -11422,6 +11664,23 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/merge-refs": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/merge-refs/-/merge-refs-2.0.0.tgz", + "integrity": "sha512-3+B21mYK2IqUWnd2EivABLT7ueDhb0b8/dGK8LoFQPrU61YITeCMn14F7y7qZafWNZhUEKb24cJdiT5Wxs3prg==", + "license": "MIT", + "funding": { + "url": "https://github.com/wojtekmaj/merge-refs?sponsor=1" + }, + "peerDependencies": { + "@types/react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + } + } + }, "node_modules/merge-stream": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/merge-stream/-/merge-stream-2.0.0.tgz", @@ -11438,14 +11697,14 @@ } }, "node_modules/mermaid": { - "version": "11.14.0", - "resolved": "https://registry.npmjs.org/mermaid/-/mermaid-11.14.0.tgz", - "integrity": "sha512-GSGloRsBs+JINmmhl0JDwjpuezCsHB4WGI4NASHxL3fHo3o/BRXTxhDLKnln8/Q0lRFRyDdEjmk1/d5Sn1Xz8g==", + "version": "11.15.0", + "resolved": "https://registry.npmjs.org/mermaid/-/mermaid-11.15.0.tgz", + "integrity": "sha512-pTMbcf3rWdtLiYGpmoTjHEpeY8seiy6sR+9nD7LOs8KfUbHE4lOUAprTRqRAcWSQ6MQpdX+YEsxShtGsINtPtw==", "license": "MIT", "dependencies": { "@braintree/sanitize-url": "^7.1.1", "@iconify/utils": "^3.0.2", - "@mermaid-js/parser": "^1.1.0", + "@mermaid-js/parser": "^1.1.1", "@types/d3": "^7.4.3", "@upsetjs/venn.js": "^2.0.0", "cytoscape": "^3.33.1", @@ -11456,14 +11715,14 @@ "dagre-d3-es": "7.0.14", "dayjs": "^1.11.19", "dompurify": "^3.3.1", + "es-toolkit": "^1.45.1", "katex": "^0.16.25", "khroma": "^2.1.0", - "lodash-es": "^4.17.23", "marked": "^16.3.0", "roughjs": "^4.6.6", "stylis": "^4.3.6", "ts-dedent": "^2.2.0", - "uuid": "^11.1.0" + "uuid": "^11.1.0 || ^12 || ^13 || ^14.0.0" } }, "node_modules/micromark": { @@ -12766,6 +13025,18 @@ "integrity": "sha512-WUjGcAqP1gQacoQe+OBJsFA7Ld4DyXuUIjZ5cc75cLHvJ7dtNsTugphxIADwspS+AraAUePCKrSVtPLFj/F88w==", "license": "MIT" }, + "node_modules/pdfjs-dist": { + "version": "5.4.296", + "resolved": "https://registry.npmjs.org/pdfjs-dist/-/pdfjs-dist-5.4.296.tgz", + "integrity": "sha512-DlOzet0HO7OEnmUmB6wWGJrrdvbyJKftI1bhMitK7O2N8W2gc757yyYBbINy9IDafXAV9wmKr9t7xsTaNKRG5Q==", + "license": "Apache-2.0", + "engines": { + "node": ">=20.16.0 || >=22.3.0" + }, + "optionalDependencies": { + "@napi-rs/canvas": "^0.1.80" + } + }, "node_modules/picocolors": { "version": "1.1.1", "resolved": "https://registry.npmjs.org/picocolors/-/picocolors-1.1.1.tgz", @@ -12932,9 +13203,9 @@ } }, "node_modules/qs": { - "version": "6.15.1", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.15.1.tgz", - "integrity": "sha512-6YHEFRL9mfgcAvql/XhwTvf5jKcOiiupt2FiJxHkiX1z4j7WL8J/jRHYLluORvc1XxB5rV20KoeK00gVJamspg==", + "version": "6.15.2", + "resolved": "https://registry.npmjs.org/qs/-/qs-6.15.2.tgz", + "integrity": "sha512-Rzq0KEyX/w/tEybncDgdkZrJgVUsUMk3xjh3t5bv3S1HTAtg+uOYt72+ZfwiQwKdysThkTBdL/rTi6HDmX9Ddw==", "license": "BSD-3-Clause", "dependencies": { "side-channel": "^1.1.0" @@ -13229,12 +13500,41 @@ } }, "node_modules/react-is": { - "version": "19.2.5", - "resolved": "https://registry.npmjs.org/react-is/-/react-is-19.2.5.tgz", - "integrity": "sha512-Dn0t8IQhCmeIT3wu+Apm1/YVsJXsGWi6k4sPdnBIdqMVtHtv0IGi6dcpNpNkNac0zB2uUAqNX3MHzN8c+z2rwQ==", + "version": "19.2.6", + "resolved": "https://registry.npmjs.org/react-is/-/react-is-19.2.6.tgz", + "integrity": "sha512-XjBR15BhXuylgWGuslhDKqlSayuqvqBX91BP8pauG8kd1zY8kotkNWbXksTCNRarse4kuGbe2kIY05ARtwNIvw==", "license": "MIT", "peer": true }, + "node_modules/react-pdf": { + "version": "10.4.1", + "resolved": "https://registry.npmjs.org/react-pdf/-/react-pdf-10.4.1.tgz", + "integrity": "sha512-kS/35staVCBqS29verTQJQZXw7RfsRCPO3fdJoW1KXylcv7A9dw6DZ3vJXC2w+bIBgLw5FN4pOFvKSQtkQhPfA==", + "license": "MIT", + "dependencies": { + "clsx": "^2.0.0", + "dequal": "^2.0.3", + "make-cancellable-promise": "^2.0.0", + "make-event-props": "^2.0.0", + "merge-refs": "^2.0.0", + "pdfjs-dist": "5.4.296", + "tiny-invariant": "^1.0.0", + "warning": "^4.0.0" + }, + "funding": { + "url": "https://github.com/wojtekmaj/react-pdf?sponsor=1" + }, + "peerDependencies": { + "@types/react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0", + "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0", + "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + } + } + }, "node_modules/react-redux": { "version": "9.2.0", "resolved": "https://registry.npmjs.org/react-redux/-/react-redux-9.2.0.tgz", @@ -15292,55 +15592,15 @@ "node": "^10 || ^12 || >=14" } }, - "node_modules/vscode-jsonrpc": { - "version": "8.2.0", - "resolved": "https://registry.npmjs.org/vscode-jsonrpc/-/vscode-jsonrpc-8.2.0.tgz", - "integrity": "sha512-C+r0eKJUIfiDIfwJhria30+TYWPtuHJXHtI7J0YlOmKAo7ogxP20T0zxB7HZQIFhIyvoBPwWskjxrvAtfjyZfA==", - "license": "MIT", - "engines": { - "node": ">=14.0.0" - } - }, - "node_modules/vscode-languageserver": { - "version": "9.0.1", - "resolved": "https://registry.npmjs.org/vscode-languageserver/-/vscode-languageserver-9.0.1.tgz", - "integrity": "sha512-woByF3PDpkHFUreUa7Hos7+pUWdeWMXRd26+ZX2A8cFx6v/JPTtd4/uN0/jB6XQHYaOlHbio03NTHCqrgG5n7g==", + "node_modules/warning": { + "version": "4.0.3", + "resolved": "https://registry.npmjs.org/warning/-/warning-4.0.3.tgz", + "integrity": "sha512-rpJyN222KWIvHJ/F53XSZv0Zl/accqHR8et1kpaMTD/fLCRxtV8iX8czMzY7sVZupTI3zcUTg8eycS2kNF9l6w==", "license": "MIT", "dependencies": { - "vscode-languageserver-protocol": "3.17.5" - }, - "bin": { - "installServerIntoExtension": "bin/installServerIntoExtension" + "loose-envify": "^1.0.0" } }, - "node_modules/vscode-languageserver-protocol": { - "version": "3.17.5", - "resolved": "https://registry.npmjs.org/vscode-languageserver-protocol/-/vscode-languageserver-protocol-3.17.5.tgz", - "integrity": "sha512-mb1bvRJN8SVznADSGWM9u/b07H7Ecg0I3OgXDuLdn307rl/J3A9YD6/eYOssqhecL27hK1IPZAsaqh00i/Jljg==", - "license": "MIT", - "dependencies": { - "vscode-jsonrpc": "8.2.0", - "vscode-languageserver-types": "3.17.5" - } - }, - "node_modules/vscode-languageserver-textdocument": { - "version": "1.0.12", - "resolved": "https://registry.npmjs.org/vscode-languageserver-textdocument/-/vscode-languageserver-textdocument-1.0.12.tgz", - "integrity": "sha512-cxWNPesCnQCcMPeenjKKsOCKQZ/L6Tv19DTRIGuLWe32lyzWhihGVJ/rcckZXJxfdKCFvRLS3fpBIsV/ZGX4zA==", - "license": "MIT" - }, - "node_modules/vscode-languageserver-types": { - "version": "3.17.5", - "resolved": "https://registry.npmjs.org/vscode-languageserver-types/-/vscode-languageserver-types-3.17.5.tgz", - "integrity": "sha512-Ld1VelNuX9pdF39h2Hgaeb5hEZM2Z3jUrrMgWQAu82jMtZp7p3vJT3BzToKtZI7NgQssZje5o0zryOrhQvzQAg==", - "license": "MIT" - }, - "node_modules/vscode-uri": { - "version": "3.1.0", - "resolved": "https://registry.npmjs.org/vscode-uri/-/vscode-uri-3.1.0.tgz", - "integrity": "sha512-/BpdSx+yCQGnCvecbyXdxHDkuk55/G3xwnC0GqY4gmQ3j+A+g8kzzgB4Nk/SINjqn6+waqw3EgbVF2QKExkRxQ==", - "license": "MIT" - }, "node_modules/web-namespaces": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/web-namespaces/-/web-namespaces-2.0.1.tgz", diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 83b1fd96f9..c49b62ab50 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -12,6 +12,7 @@ "lint": "eslint .", "preview": "vite preview", "typecheck": "tsc -b --pretty false", + "i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts", "biome:check": "biome check", "biome:fix": "biome check --write" }, @@ -39,12 +40,14 @@ "@tailwindcss/vite": "^4.2.2", "@tanstack/react-router": "1.169.2", "@tanstack/react-table": "^8.21.3", + "@tanstack/react-virtual": "3.13.25", "@tauri-apps/api": "^2.10.1", "@tauri-apps/plugin-clipboard-manager": "^2.3.2", "@tauri-apps/plugin-notification": "^2.3.3", "@tauri-apps/plugin-opener": "^2.5.3", "@tauri-apps/plugin-process": "^2.3.1", "@tauri-apps/plugin-updater": "^2.10.1", + "@tauri-apps/plugin-window-state": "^2.4.1", "@toolwind/corner-shape": "^0.0.8-3", "@xyflow/react": "^12.10.0", "assistant-stream": "0.3.12", @@ -65,6 +68,7 @@ "react": "^19.2.4", "react-day-picker": "^9.13.2", "react-dom": "^19.2.4", + "react-pdf": "10.4.1", "react-resizable-panels": "^4.6.4", "recharts": "3.7.0", "shadcn": "^4.2.0", @@ -80,15 +84,20 @@ "overrides": { "@tanstack/react-router": "1.169.2", "@tanstack/router-core": "1.169.2", - "@tanstack/history": "1.161.6" + "@tanstack/history": "1.161.6", + "mermaid": "11.15.0", + "hono": "4.12.21", + "qs": "6.15.2", + "ip-address": "10.1.1", + "brace-expansion@5.0.5": "5.0.6" }, "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", + "@types/node-forge": "^1.3.14", "@types/react": "^19.2.5", "@types/react-dom": "^19.2.3", "@vitejs/plugin-react": "^6.0.1", @@ -99,5 +108,10 @@ "typescript": "~5.9.3", "typescript-eslint": "^8.55.0", "vite": "^8.0.1" + }, + "allowScripts": { + "@biomejs/biome@1.9.4": true, + "msw@2.14.3": true, + "fsevents": true } } diff --git a/studio/frontend/public/hub/profile/logo/anthropic.svg b/studio/frontend/public/hub/profile/logo/anthropic.svg new file mode 100644 index 0000000000..7545cc8f3e --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/anthropic.svg @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/cohere.png b/studio/frontend/public/hub/profile/logo/cohere.png new file mode 100644 index 0000000000..99eabbb54f Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/cohere.png differ diff --git a/studio/frontend/public/hub/profile/logo/deepseek.svg b/studio/frontend/public/hub/profile/logo/deepseek.svg new file mode 100644 index 0000000000..d1ba06b942 --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/deepseek.svg @@ -0,0 +1,14 @@ + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/google.png b/studio/frontend/public/hub/profile/logo/google.png new file mode 100644 index 0000000000..01bb81206e Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/google.png differ diff --git a/studio/frontend/public/hub/profile/logo/hf.svg b/studio/frontend/public/hub/profile/logo/hf.svg new file mode 100644 index 0000000000..ab959d165f --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/hf.svg @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/studio/frontend/public/hub/profile/logo/ibm.png b/studio/frontend/public/hub/profile/logo/ibm.png new file mode 100644 index 0000000000..31c965f0b3 Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/ibm.png differ diff --git a/studio/frontend/public/hub/profile/logo/meta.svg b/studio/frontend/public/hub/profile/logo/meta.svg new file mode 100644 index 0000000000..9fa656bd6b --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/meta.svg @@ -0,0 +1,19 @@ + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/microsoft.svg b/studio/frontend/public/hub/profile/logo/microsoft.svg new file mode 100644 index 0000000000..5334aa7ca6 --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/microsoft.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/minimax-color.png b/studio/frontend/public/hub/profile/logo/minimax-color.png new file mode 100644 index 0000000000..e9472c676d Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/minimax-color.png differ diff --git a/studio/frontend/public/hub/profile/logo/mistral.svg b/studio/frontend/public/hub/profile/logo/mistral.svg new file mode 100644 index 0000000000..40c2591b31 --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/mistral.svg @@ -0,0 +1,19 @@ + + + + + + + + + + + + + + + + + + + diff --git a/studio/frontend/public/hub/profile/logo/moonshot.jpg b/studio/frontend/public/hub/profile/logo/moonshot.jpg new file mode 100644 index 0000000000..956a5b58b1 Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/moonshot.jpg differ diff --git a/studio/frontend/public/hub/profile/logo/nvidia.svg b/studio/frontend/public/hub/profile/logo/nvidia.svg new file mode 100644 index 0000000000..ae65b09a2b --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/nvidia.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/openai.svg b/studio/frontend/public/hub/profile/logo/openai.svg new file mode 100644 index 0000000000..74d9b1b44b --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/openai.svg @@ -0,0 +1,5 @@ + + + + + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/qwen.png b/studio/frontend/public/hub/profile/logo/qwen.png new file mode 100644 index 0000000000..67d2258f40 Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/qwen.png differ diff --git a/studio/frontend/public/hub/profile/logo/xai.svg b/studio/frontend/public/hub/profile/logo/xai.svg new file mode 100644 index 0000000000..0c83eb3d9b --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/xai.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/zai.svg b/studio/frontend/public/hub/profile/logo/zai.svg new file mode 100644 index 0000000000..28ca7280a1 --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/zai.svg @@ -0,0 +1,215 @@ + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 83238dadf0..91d99e238e 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -9,7 +9,10 @@ import { shouldUseCustomWindowTitlebar, } from "@/components/tauri/window-titlebar"; import { Toaster } from "@/components/ui/sonner"; +import { TooltipProvider } from "@/components/ui/tooltip"; import { WebUpdateBanner } from "@/components/web/update-banner"; +import { LlamaUpdateBanner } from "@/components/llama-update-banner"; +import { DownloadManagerPanel } from "@/features/hub/download-manager"; import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth"; import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain"; import { useTauriBackend, type BackendStatus } from "@/hooks/use-tauri-backend"; @@ -26,6 +29,9 @@ interface AppProviderProps { type TauriWindowMode = "setup" | "app"; type WindowLayoutGuard = () => boolean; +const MIN_WINDOW_WIDTH = 900; +const MIN_WINDOW_HEIGHT = 600; + async function showSetupWindow(isCurrent: WindowLayoutGuard): Promise { const { getCurrentWindow } = await import("@tauri-apps/api/window"); if (!isCurrent()) return; @@ -39,35 +45,52 @@ async function showSetupWindow(isCurrent: WindowLayoutGuard): Promise { async function applyAppWindowLayout(isCurrent: WindowLayoutGuard): Promise { const { getCurrentWindow, currentMonitor, LogicalSize } = await import("@tauri-apps/api/window"); + const { invoke } = await import("@tauri-apps/api/core"); + const { restoreStateCurrent, StateFlags } = await import("@tauri-apps/plugin-window-state"); if (!isCurrent()) return; const win = getCurrentWindow(); - const monitor = await currentMonitor(); + // Decide first-launch vs restore from the on-disk state file BEFORE touching the + // window. Probing the window after restoreStateCurrent is unreliable: on GTK, + // set_size on a hidden window is deferred until show(), so innerSize() reads a + // stale value and a baseline fallback would overwrite the queued restore. On + // macOS the same probe works, hence the inconsistency between prior iterations. + const hasSavedState = await invoke("has_saved_window_state"); if (!isCurrent()) return; - let finalW = 900; - let finalH = 600; - - if (monitor) { - const scale = monitor.scaleFactor; - const screenW = monitor.size.width / scale; - const screenH = monitor.size.height / scale; - - finalW = Math.max(900, Math.round(screenW * 0.75)); - const targetH = Math.max(600, Math.round(finalW / 1.618)); - finalH = Math.min(targetH, Math.round(screenH * 0.85)); - } - - if (!isCurrent()) return; - await win.setSize(new LogicalSize(finalW, finalH)); - if (!isCurrent()) return; - await win.setSizeConstraints({ minWidth: 900, minHeight: 600 }); - if (!isCurrent()) return; await win.setResizable(true); if (!isCurrent()) return; - await win.center(); + + if (hasSavedState) { + // Subsequent launch: plugin restores size/position/maximized, with built-in + // off-screen protection for positions saved on a now-disconnected display. + await restoreStateCurrent( + StateFlags.SIZE | StateFlags.POSITION | StateFlags.MAXIMIZED, + ); + } else { + // First launch: fit to the current monitor and center. + const monitor = await currentMonitor(); + if (!isCurrent()) return; + let finalW = MIN_WINDOW_WIDTH; + let finalH = MIN_WINDOW_HEIGHT; + if (monitor) { + const scale = monitor.scaleFactor; + const screenW = monitor.size.width / scale; + const screenH = monitor.size.height / scale; + finalW = Math.max(MIN_WINDOW_WIDTH, Math.round(screenW * 0.75)); + const targetH = Math.max(MIN_WINDOW_HEIGHT, Math.round(finalW / 1.618)); + finalH = Math.min(targetH, Math.round(screenH * 0.85)); + } + await win.setSize(new LogicalSize(finalW, finalH)); + if (!isCurrent()) return; + await win.center(); + } if (!isCurrent()) return; await win.show(); + if (!isCurrent()) return; + // Apply constraints after restore/show: doing so before plugin restore can emit + // a Resized event and overwrite the plugin's cached saved size. + await win.setSizeConstraints({ minWidth: MIN_WINDOW_WIDTH, minHeight: MIN_WINDOW_HEIGHT }); } async function showWindowFallback(): Promise { @@ -235,7 +258,9 @@ function TauriWrapper({ children }: { children: ReactNode }) { return ( <> {children} + + ); } @@ -252,6 +277,7 @@ function TauriWrapper({ children }: { children: ReactNode }) { {children} + ) : ( ); - if (!shouldUseCustomWindowTitlebar()) return content; + if (!shouldUseCustomWindowTitlebar()) { + // macOS desktop uses the native titlebar and returns here before the + // custom-titlebar branch, so mount the updater banner on this path too. + return ( + <> + {content} + + + ); + } const showSidebarSurface = showApp && !HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname); @@ -281,6 +318,9 @@ function TauriWrapper({ children }: { children: ReactNode }) {
{content}
+ ); } @@ -288,17 +328,19 @@ function TauriWrapper({ children }: { children: ReactNode }) { export function AppProvider({ children }: AppProviderProps) { return ( - - {children} - - + + + {children} + + + ); } diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index c7bc0440bd..5c18e637e2 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -3,6 +3,8 @@ import { Link, createRouter, useRouterState } from "@tanstack/react-router"; import { Button } from "@/components/ui/button"; +import { MascotImg } from "@/components/mascot-img"; +import { useT } from "@/i18n"; import { Route as rootRoute } from "./routes/__root"; import { Route as dataRecipesRoute } from "./routes/data-recipes"; import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; @@ -11,7 +13,9 @@ import { Route as exportRoute } from "./routes/export"; import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; import { Route as loginRoute } from "./routes/login"; +import { Route as hubRoute } from "./routes/hub"; import { Route as onboardingRoute } from "./routes/onboarding"; +import { Route as projectsRoute } from "./routes/projects"; import { Route as changePasswordRoute } from "./routes/change-password"; import { Route as settingsRoute } from "./routes/settings"; import { Route as studioRoute } from "./routes/studio"; @@ -22,33 +26,33 @@ const routeTree = rootRoute.addChildren([ loginRoute, changePasswordRoute, gridTestRoute, + hubRoute, settingsRoute, studioRoute, chatRoute, + projectsRoute, exportRoute, dataRecipesRoute, dataRecipeRoute, ]); function DefaultNotFound() { + const t = useT(); const pathname = useRouterState({ select: (s) => s.location.pathname }); + return (
- Sloth mascot +

- Page not found + {t("shell.notFound.title")}

- {pathname} does not exist. + {t("shell.notFound.description", { path: pathname })}

); diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 47bff815e6..c05b0e6451 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -6,8 +6,10 @@ import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; -import { useTrainingUnloadGuard } from "@/features/training/hooks/use-training-unload-guard"; +import { useChatRuntimeStore } from "@/features/chat"; +import { useTrainingUnloadGuard } from "@/features/training"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; +import { useT, type TranslationKey } from "@/i18n"; import { Outlet, createRootRoute, @@ -16,28 +18,31 @@ import { useRouterState, } from "@tanstack/react-router"; import { AnimatePresence, motion } from "motion/react"; -import { Suspense, useEffect, useLayoutEffect, type ReactNode } from "react"; +import { Suspense, useEffect, useLayoutEffect } from "react"; import { AppProvider } from "../provider"; -// Type `staticData.title` on every route so the matched-title selector -// below stays type-safe without an inline cast. declare module "@tanstack/react-router" { interface StaticDataRouteOption { title?: string; + titleKey?: TranslationKey; } } -// Fallback while a lazy route bundle (Train/Recipes/Export) loads. -// /chat is synchronous and never hits this. -const RouteFallback: ReactNode = ( -
- Loading... -
-); +function RouteFallback() { + const t = useT(); + + return ( +
+ {t("common.loading")} +
+ ); +} const CHAT_ONLY_ALLOWED = new Set([ "/", "/chat", + "/projects", + "/hub", "/login", "/signup", "/change-password", @@ -51,8 +56,8 @@ function isChatOnlyAllowed(pathname: string): boolean { export const Route = createRootRoute({ beforeLoad: async ({ location }) => { - // Ensure platform info is fetched before checking chat-only guard. - // fetchDeviceType caches after first call, so subsequent navigations are instant. + // Fetch platform info before the chat-only guard. fetchDeviceType caches, + // so later navigations are instant. await fetchDeviceType(); const chatOnly = usePlatformStore.getState().isChatOnly(); if (chatOnly && !isChatOnlyAllowed(location.pathname)) { @@ -68,6 +73,7 @@ const HIDDEN_NAVBAR_ROUTES = ["/onboarding", "/login", "/change-password"]; const DEFAULT_DOCUMENT_TITLE = "Unsloth Studio"; function RootLayout() { + const t = useT(); const pathname = useRouterState({ select: (s) => s.location.pathname }); const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname); const isChatRoute = pathname.startsWith("/chat"); @@ -75,24 +81,20 @@ function RootLayout() { useTrainingUnloadGuard(); - // Walk matches deepest-first; each route declares its own title. const matchedTitle = useMatches({ select: (matches) => { for (let i = matches.length - 1; i >= 0; i--) { - const title = matches[i].staticData.title; + const { title, titleKey } = matches[i].staticData; + if (titleKey) return t(titleKey); if (title) return title; } return null; }, }); - // `/settings` redirects in `beforeLoad`, so its route never stays - // matched; surface the modal's title via the store instead. const settingsDialogOpen = useSettingsDialogStore((s) => s.open); - const documentTitle = settingsDialogOpen ? "Settings" : matchedTitle; + const documentTitle = settingsDialogOpen ? t("settings.title") : matchedTitle; - // useLayoutEffect updates the tab title before paint, avoiding a - // one-frame flash of the previous route's title on navigation. useLayoutEffect(() => { document.title = documentTitle ? `${documentTitle} - ${DEFAULT_DOCUMENT_TITLE}` @@ -111,12 +113,19 @@ function RootLayout() { return () => window.removeEventListener("keydown", handler); }, []); + useEffect(() => { + if (isChatRoute) return; + const chatRuntime = useChatRuntimeStore.getState(); + chatRuntime.setActiveProjectId(null); + chatRuntime.setActiveThreadId(null); + }, [isChatRoute]); + return ( {hideNavbar ? (
- + }>
@@ -131,9 +140,14 @@ function RootLayout() {
- + {/* Use mode="popLayout" instead of "wait" to prevent UI freezes when + switching from heavy pages (like Export with many checkpoints). + "popLayout" allows the new route to mount immediately while the + old one animates out, avoiding blocking on expensive exit renders. + See issue #5850. */} + - + }> diff --git a/studio/frontend/src/app/routes/chat.tsx b/studio/frontend/src/app/routes/chat.tsx index 98c73aa7e0..a5514cdee1 100644 --- a/studio/frontend/src/app/routes/chat.tsx +++ b/studio/frontend/src/app/routes/chat.tsx @@ -10,6 +10,7 @@ export type ChatSearch = { thread?: string; compare?: string; new?: string; + project?: string; }; export const Route = createRoute({ @@ -21,6 +22,7 @@ export const Route = createRoute({ thread: typeof search.thread === "string" ? search.thread : undefined, compare: typeof search.compare === "string" ? search.compare : undefined, new: typeof search.new === "string" ? search.new : undefined, + project: typeof search.project === "string" ? search.project : undefined, }), component: ChatPage, }); diff --git a/studio/frontend/src/app/routes/hub.tsx b/studio/frontend/src/app/routes/hub.tsx new file mode 100644 index 0000000000..dcd6617ec8 --- /dev/null +++ b/studio/frontend/src/app/routes/hub.tsx @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { createRoute } from "@tanstack/react-router"; +import { lazy } from "react"; +import { requireAuth } from "../auth-guards"; +import { Route as rootRoute } from "./__root"; + +const ModelsPage = lazy(() => + import("@/features/hub/hub-page").then((m) => ({ + default: m.ModelsPage, + })), +); + +export interface ModelsSearch { + tab?: "discover" | "downloaded"; +} + +export const Route = createRoute({ + getParentRoute: () => rootRoute, + path: "/hub", + beforeLoad: () => requireAuth(), + component: ModelsPage, + validateSearch: (search: Record): ModelsSearch => { + const raw = search.tab; + if (raw === "discover" || raw === "downloaded") return { tab: raw }; + return {}; + }, +}); diff --git a/studio/frontend/src/app/routes/projects.tsx b/studio/frontend/src/app/routes/projects.tsx new file mode 100644 index 0000000000..c63b1d5838 --- /dev/null +++ b/studio/frontend/src/app/routes/projects.tsx @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { createRoute } from "@tanstack/react-router"; +import { lazy } from "react"; +import { requireAuth } from "../auth-guards"; +import { Route as rootRoute } from "./__root"; + +const ProjectsPage = lazy(() => + import("@/features/chat/projects-page").then((m) => ({ + default: m.ProjectsPage, + })), +); + +export const Route = createRoute({ + getParentRoute: () => rootRoute, + path: "/projects", + staticData: { title: "Projects" }, + beforeLoad: () => requireAuth(), + component: ProjectsPage, +}); diff --git a/studio/frontend/src/app/routes/settings.tsx b/studio/frontend/src/app/routes/settings.tsx index fa97a450f7..84650bdb8b 100644 --- a/studio/frontend/src/app/routes/settings.tsx +++ b/studio/frontend/src/app/routes/settings.tsx @@ -7,10 +7,10 @@ import { useSettingsDialogStore } from "@/features/settings"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -// /settings is a deep link to the modal. Open it, then redirect home. -// Tab title is driven by useSettingsDialogStore in __root.tsx since the -// redirect means /settings never stays matched; staticData is just a -// safety net if beforeLoad ever stops throwing. +// /settings deep-links the modal: open it, then redirect home. Tab title is +// driven by useSettingsDialogStore in __root.tsx since the redirect means +// /settings never stays matched; staticData is a safety net if beforeLoad +// ever stops throwing. export const Route = createRoute({ getParentRoute: () => rootRoute, path: "/settings", diff --git a/studio/frontend/src/app/routes/studio.tsx b/studio/frontend/src/app/routes/studio.tsx index 75f1a1b937..ae7f445e94 100644 --- a/studio/frontend/src/app/routes/studio.tsx +++ b/studio/frontend/src/app/routes/studio.tsx @@ -15,7 +15,7 @@ const StudioPage = lazy(() => export const Route = createRoute({ getParentRoute: () => rootRoute, path: "/studio", - staticData: { title: "Train" }, + staticData: { titleKey: "studio.routeTitle" }, beforeLoad: () => requireAuth(), component: StudioPage, }); diff --git a/studio/frontend/src/asset-queries.d.ts b/studio/frontend/src/asset-queries.d.ts new file mode 100644 index 0000000000..2c7a1cd936 --- /dev/null +++ b/studio/frontend/src/asset-queries.d.ts @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Vite `?inline` imports (data URIs); vite/client only types bare extensions. +declare module "*?inline" { + const src: string; + // biome-ignore lint/style/noDefaultExport: Vite asset modules export default. + export default src; +} diff --git a/studio/frontend/src/assets/mascot-fallback.webp b/studio/frontend/src/assets/mascot-fallback.webp new file mode 100644 index 0000000000..ad50361dbb Binary files /dev/null and b/studio/frontend/src/assets/mascot-fallback.webp differ diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index aac5f8f8a8..f1a200b09b 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -26,6 +26,9 @@ import { DropdownMenuItem, DropdownMenuSeparator, DropdownMenuShortcut, + DropdownMenuSub, + DropdownMenuSubContent, + DropdownMenuSubTrigger, DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { @@ -38,19 +41,26 @@ import { } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; import { Button } from "@/components/ui/button"; +import { Switch } from "@/components/ui/switch"; import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler"; import { cn } from "@/lib/utils"; import { ChefHatIcon, - ColumnInsertIcon, CursorInfo02Icon, + DashboardCircleIcon, Delete02Icon, + Download01Icon, DownloadSquare01Icon, Edit03Icon, + FolderAddIcon, + FolderExportIcon, + Folder01Icon, Globe02Icon, HelpCircleIcon, - Logout01Icon, + Logout05Icon, + MoreVerticalIcon, Search01Icon, + PlusSignIcon, PowerIcon, PencilEdit02Icon, LayoutAlignLeftIcon, @@ -58,6 +68,12 @@ import { TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; +import { + exportConversationRawJsonl, + exportConversationCsv, + exportConversationShareGPT, +} from "@/features/chat/prompt-storage/prompt-storage-dialog"; +import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, TooltipContent, @@ -68,11 +84,17 @@ import { ChevronDown, ChevronsUpDown, MoreHorizontalIcon, Moon, Sun } from "luci import { Link, useNavigate, useRouterState } from "@tanstack/react-router"; import { ChatSearchDialog, + createChatProject, + deleteChatProject, deleteChatItem, + moveChatItemToProject, renameChatItem, + renameChatProject, useChatRuntimeStore, + useChatProjects, useChatSearchStore, useChatSidebarItems, + type ProjectRecord, type SidebarItem, } from "@/features/chat"; import { useSettingsDialogStore } from "@/features/settings"; @@ -90,9 +112,33 @@ import { useTrainingRuntimeStore, } from "@/features/training"; import type { TrainingRunSummary } from "@/features/training"; -import { useEffect, useRef, useState } from "react"; +import { useEffect, useMemo, useRef, useState, type ReactNode } from "react"; import { toast } from "@/lib/toast"; import { ShutdownDialog } from "@/components/shutdown-dialog"; +import { translate, useT, type TranslationKey } from "@/i18n"; + +const EMPHASIS_MARKER = "__UNSLOTH_I18N_EMPHASIS_MARKER__"; + +type AppT = ReturnType; + +function renderEmphasizedTranslation( + t: AppT, + key: TranslationKey, + emphasizedValue: string, +): ReactNode { + const translated = t(key, { name: EMPHASIS_MARKER }); + const parts = translated.split(EMPHASIS_MARKER); + if (parts.length === 1) return translated; + + const nodes: ReactNode[] = []; + parts.forEach((part, index) => { + if (part.length > 0) nodes.push(part); + if (index < parts.length - 1) { + nodes.push({emphasizedValue}); + } + }); + return nodes; +} function getTourId(pathname: string): string | null { if (pathname.startsWith("/studio")) return "studio"; @@ -101,11 +147,8 @@ function getTourId(pathname: string): string | null { return null; } -// Hugeicons' TestTube01Icon ships with two interior bubbles (paths #4 -// and #5 of the 5-path definition). Slicing to the first three paths -// keeps the test-tube outline + horizontal cap + liquid line, dropping -// the bubbles. The original export stays untouched, and HugeiconsIcon -// renders this trimmed array exactly the same way. +// TestTube01Icon's last 2 paths are interior bubbles; slice to the first +// 3 (outline + cap + liquid line) to drop them. Original export untouched. const TestTubeOutlineIcon = TestTube01Icon.slice( 0, 3, @@ -155,17 +198,19 @@ function NavItem({ onClick, children, dataTour, + className, }: { icon: typeof ZapIcon; label: string; active: boolean; disabled?: boolean; onClick: () => void; - children?: React.ReactNode; + children?: ReactNode; dataTour?: string; + className?: string; }) { return ( - +
{label} @@ -185,6 +230,7 @@ function NavItem({ } export function AppSidebar() { + const t = useT(); const { isDark, toggleTheme, anchorRef } = useAnimatedThemeToggle(); const { pathname, search } = useRouterState({ select: (s) => ({ @@ -204,30 +250,56 @@ export function AppSidebar() { const chatOnly = usePlatformStore((s) => s.isChatOnly()); const [shutdownOpen, setShutdownOpen] = useState(false); - // Chat collapsible state — open by default, auto-expand on route entry const isChatRoute = pathname.startsWith("/chat"); const isStudioRoute = pathname === "/studio" || pathname.startsWith("/studio/"); const [chatOpen, setChatOpen] = useState(true); + + const [trainOpen, setTrainOpen] = useState(true); const [runsOpen, setRunsOpen] = useState(true); - useEffect(() => { if (isChatRoute) setChatOpen(true); }, [isChatRoute]); - useEffect(() => { if (isStudioRoute) setRunsOpen(true); }, [isStudioRoute]); + useEffect(() => { + if (!isChatRoute) return; + queueMicrotask(() => setChatOpen(true)); + }, [isChatRoute]); + useEffect(() => { + if (!isStudioRoute) return; + queueMicrotask(() => setRunsOpen(true)); + }, [isStudioRoute]); const scrollRef = useRef(null); const [scrolled, setScrolled] = useState(false); - useEffect(() => { - const el = scrollRef.current; - if (!el) return; - const handler = () => setScrolled(el.scrollTop > 0); - handler(); - el.addEventListener("scroll", handler, { passive: true }); - return () => el.removeEventListener("scroll", handler); - }, []); + // Bottom fade hides at the very bottom / for short lists so the last row + // isn't washed out (Gemini-style). + const [canScrollDown, setCanScrollDown] = useState(false); + // Driven only from onScroll + a content-change effect below. No + // ResizeObserver: its callback-driven setState caused a render loop (React + // #185). Both setters bail out when unchanged, so neither path can loop. + const syncScrollState = (el: HTMLDivElement) => { + const nextScrolled = el.scrollTop > 0; + setScrolled((prev) => (prev === nextScrolled ? prev : nextScrolled)); + const nextCanScrollDown = + el.scrollHeight - el.scrollTop - el.clientHeight > 1; + setCanScrollDown((prev) => + prev === nextCanScrollDown ? prev : nextCanScrollDown, + ); + }; const isRecipesRoute = pathname.startsWith("/data-recipes"); const { displayTitle, avatarDataUrl } = useEffectiveProfile(); - const { items: chatItems } = useChatSidebarItems(); + const { projects } = useChatProjects(); + const activeProjectId = isChatRoute + ? ((search.project as string | undefined) ?? null) + : null; + const { items: allChatItems } = useChatSidebarItems({ + enabled: !isStudioRoute, + requireMessages: false, + }); + const recentChatItems = useMemo( + () => allChatItems.filter((item) => !item.projectId), + [allChatItems], + ); + const chatItems = allChatItems; const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId); const activeThreadId = isChatRoute @@ -242,33 +314,86 @@ export function AppSidebar() { !chatOnly && isStudioRoute, ); const activeJobId = useTrainingRuntimeStore((s) => s.jobId); + const currentRunViewActive = useTrainingRuntimeStore((s) => s.currentRunViewActive); const selectedHistoryRunId = useTrainingRuntimeStore((s) => s.selectedHistoryRunId); const setSelectedHistoryRunId = useTrainingRuntimeStore((s) => s.setSelectedHistoryRunId); + // Recompute bottom-fade on mount and whenever list height can change + // (items load, sections toggle, route switch) - onScroll never fires for + // short, non-scrolling lists. Guarded setState below can't loop. + useEffect(() => { + const el = scrollRef.current; + if (!el) return; + const next = el.scrollHeight - el.scrollTop - el.clientHeight > 1; + setCanScrollDown((prev) => (prev === next ? prev : next)); + }, [ + recentChatItems.length, + runItems.length, + projects.length, + chatOpen, + trainOpen, + runsOpen, + isStudioRoute, + ]); + const chatDisabled = isTrainingRunning; + function chatSearchForProject(projectId: string | null) { + if (projectId) { + return { project: projectId }; + } + return { + new: createNavigationNonce(), + }; + } + + function openNewChat(projectId = activeProjectId) { + if (chatDisabled) return; + setActiveThreadId(null); + useChatRuntimeStore.getState().setActiveProjectId(projectId); + navigate({ to: "/chat", search: chatSearchForProject(projectId) }); + closeMobileIfOpen(); + } + + function openProject(projectId: string) { + if (chatDisabled) return; + setActiveThreadId(null); + useChatRuntimeStore.getState().setActiveProjectId(projectId); + navigate({ to: "/chat", search: { project: projectId } }); + closeMobileIfOpen(); + } + async function handleDeleteThread(item: Parameters[0]) { await deleteChatItem(item, activeThreadId, (view) => { navigate({ to: "/chat", - search: { new: view.newThreadNonce }, + search: item.projectId + ? { project: item.projectId } + : { new: view.newThreadNonce }, }); }); } type RenameTarget = | { kind: "chat"; item: SidebarItem; current: string } + | { kind: "project"; project: ProjectRecord; current: string } | { kind: "run"; run: TrainingRunSummary; current: string }; const [renamingTarget, setRenamingTarget] = useState( null, ); const [renameDraft, setRenameDraft] = useState(""); + const [creatingProject, setCreatingProject] = useState(false); + const [projectNameDraft, setProjectNameDraft] = useState(""); + const [projectCreateMoveTarget, setProjectCreateMoveTarget] = + useState(null); const renameTrimmed = renameDraft.trim(); const nextRunDisplayName = renameTrimmed.length > 0 ? renameTrimmed : null; const renameDirty = renamingTarget !== null && (renamingTarget.kind === "chat" ? renameTrimmed.length > 0 && renameTrimmed !== renamingTarget.current + : renamingTarget.kind === "project" + ? renameTrimmed.length > 0 && renameTrimmed !== renamingTarget.current : renameTrimmed.length > 0 ? renameTrimmed !== renamingTarget.current : renamingTarget.run.display_name != null); @@ -290,7 +415,17 @@ export function AppSidebar() { try { await renameChatItem(target.item, renameTrimmed); } catch (err) { - toast.error("Failed to rename chat", { + toast.error(translate("shell.toast.failedToRenameChat"), { + description: err instanceof Error ? err.message : undefined, + }); + } + return; + } + if (target.kind === "project") { + try { + await renameChatProject(target.project.id, renameTrimmed); + } catch (err) { + toast.error("Failed to rename project", { description: err instanceof Error ? err.message : undefined, }); } @@ -300,7 +435,7 @@ export function AppSidebar() { const updated = await renameTrainingRun(target.run.id, nextRunDisplayName); emitTrainingRunUpdated(updated); } catch (err) { - toast.error("Failed to rename run", { + toast.error(translate("shell.toast.failedToRenameRun"), { description: err instanceof Error ? err.message : undefined, }); } @@ -308,26 +443,52 @@ export function AppSidebar() { type DeleteTarget = | { kind: "chat"; item: SidebarItem } + | { kind: "project"; project: ProjectRecord } | { kind: "run"; run: TrainingRunSummary }; const [confirmingDelete, setConfirmingDelete] = useState(null); + const [deleteProjectFiles, setDeleteProjectFiles] = useState(false); + + useEffect(() => { + if (confirmingDelete?.kind !== "project") { + setDeleteProjectFiles(false); + } + }, [confirmingDelete]); async function commitDelete() { const target = confirmingDelete; if (!target) return; + const shouldDeleteProjectFiles = + target.kind === "project" && deleteProjectFiles; setConfirmingDelete(null); if (target.kind === "chat") { try { await handleDeleteThread(target.item); } catch (err) { - toast.error("Failed to delete chat", { + toast.error(translate("shell.toast.failedToDeleteChat"), { + description: err instanceof Error ? err.message : undefined, + }); + } + return; + } + if (target.kind === "project") { + try { + await deleteChatProject(target.project.id, { + deleteFiles: shouldDeleteProjectFiles, + }); + if (activeProjectId === target.project.id) { + useChatRuntimeStore.getState().setActiveProjectId(null); + navigate({ to: "/chat", search: { new: createNavigationNonce() } }); + } + } catch (err) { + toast.error("Failed to delete project", { description: err instanceof Error ? err.message : undefined, }); } return; } if (target.run.status === "running") { - toast.error("Cannot delete a running training run"); + toast.error(t("shell.toast.cannotDeleteRunningRun")); return; } try { @@ -337,12 +498,213 @@ export function AppSidebar() { } emitTrainingRunDeleted(target.run.id); } catch (err) { - toast.error("Failed to delete run", { + toast.error(translate("shell.toast.failedToDeleteRun"), { description: err instanceof Error ? err.message : undefined, }); } } + async function commitCreateProject() { + const name = projectNameDraft.trim(); + if (!name) return; + const moveTarget = projectCreateMoveTarget; + try { + const project = await createChatProject(name); + if (moveTarget) { + await moveChatItemToProject(moveTarget, project.id); + if (activeThreadId === moveTarget.id) { + useChatRuntimeStore.getState().setActiveProjectId(project.id); + } + } + setCreatingProject(false); + setProjectNameDraft(""); + setProjectCreateMoveTarget(null); + if (moveTarget) { + return; + } else { + openProject(project.id); + } + } catch (err) { + toast.error(moveTarget ? "Failed to create and move chat" : "Failed to create project", { + description: err instanceof Error ? err.message : undefined, + }); + } + } + + async function moveChatToProject(item: SidebarItem, projectId: string | null) { + if (item.projectId === projectId) return; + try { + await moveChatItemToProject(item, projectId); + if (activeThreadId === item.id) { + useChatRuntimeStore.getState().setActiveProjectId(projectId); + } + } catch (err) { + toast.error("Failed to move chat", { + description: err instanceof Error ? err.message : undefined, + }); + } + } + + function renderChatSidebarItem( + item: SidebarItem, + variant: "project" | "recent", + ) { + const itemClass = + variant === "project" + ? "group/project-chat-item relative" + : "group/recent-item relative"; + const actionClass = + variant === "project" + ? "sidebar-row-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" + : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"; + const buttonClass = cn( + "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium", + // pl-3.5 starts the title at the same x as the Recents label text. + variant === "project" ? "pl-[39px]" : "pl-3.5", + variant === "project" + ? "group-hover/project-chat-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8" + : "group-hover/recent-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8", + ); + + return ( + + { + navigate({ + to: "/chat", + search: + item.type === "single" + ? { + thread: item.id, + ...(item.projectId ? { project: item.projectId } : {}), + } + : { + compare: item.id, + ...(item.projectId ? { project: item.projectId } : {}), + }, + }); + closeMobileIfOpen(); + }} + > + {item.title} + + + + + + + openRenameChat(item)}> + + Rename + + + + + Move to project + + + { + setProjectCreateMoveTarget(item); + setProjectNameDraft(""); + setCreatingProject(true); + }} + > + + New project + + void moveChatToProject(item, null)} + > + Recents + + {projects.map((project) => ( + void moveChatToProject(item, project.id)} + > + + {project.name} + + ))} + + + + + + Export + + + {[ + { label: "Raw JSONL", fn: exportConversationRawJsonl }, + { label: "CSV", fn: exportConversationCsv }, + { label: "ShareGPT JSONL", fn: exportConversationShareGPT }, + ].map(({ label, fn }) => ( + { + try { + const ids = item.type === "single" + ? [item.id] + : (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id); + await Promise.all(ids.map((id) => fn(id))); + } catch { + toast.error("Export failed."); + } + }} + > + {label} + + ))} + + {/* Bulk export and import live in Settings -> Chat -> Data. */} + + useSettingsDialogStore.getState().openDialog("chat") + } + > + Export all chats… + + + + setConfirmingDelete({ kind: "chat", item })} + > + + Delete + + + + + ); + } + return ( <> - + {/* Expanded: compact logo + close toggle */}
{ event.preventDefault(); if (chatDisabled) return; - setActiveThreadId(null); - closeMobileIfOpen(); - void navigate({ - to: "/chat", - search: { new: createNavigationNonce() }, - }); + openNewChat(null); }} className="flex items-center gap-[6px] select-none" - aria-label="Unsloth home" + aria-label={t("shell.aria.home")} > Unsloth - + unsloth - BETA + {t("shell.beta")} {!isMobile && ( @@ -386,8 +743,8 @@ export function AppSidebar() { @@ -397,7 +754,7 @@ export function AppSidebar() { sideOffset={6} className="tooltip-compact" > - Close sidebar + {t("shell.aria.closeSidebar")} )} @@ -405,14 +762,14 @@ export function AppSidebar() { {/* Collapsed: panel icon doubles as expand trigger */} {!isMobile && ( -
+
@@ -422,48 +779,35 @@ export function AppSidebar() { sideOffset={8} className="tooltip-compact" > - Open sidebar + {t("shell.aria.openSidebar")}
)} - + { - if (chatDisabled) return; - setActiveThreadId(null); - navigate({ to: "/chat", search: { new: createNavigationNonce() } }); - closeMobileIfOpen(); - }} - /> - i.id === search.compare)} - disabled={chatDisabled} - dataTour="chat-compare" - onClick={() => { - if (chatDisabled) return; - setActiveThreadId(null); - navigate({ to: "/chat", search: { compare: createNavigationNonce() } }); - closeMobileIfOpen(); - }} + onClick={() => openNewChat(null)} /> { - if (chatDisabled) return; + // Search is read-only and never runs inference, so it stays + // available while training (unlike New chat, gated on chatDisabled). useChatSearchStore.getState().open(); closeMobileIfOpen(); }} @@ -472,138 +816,172 @@ export function AppSidebar() { - - - - { - if (chatOnly) return; - navigate({ to: "/studio" }); - closeMobileIfOpen(); - }} - /> + syncScrollState(e.currentTarget)} + className={cn( + // pb-2 keeps the last row's rounded highlight clear of the + // overflow clip edge so its bottom corners aren't shaved off. + "sidebar-scroll-fade gap-0 overflow-y-auto overscroll-contain min-h-0 pb-2", + scrolled && "is-scrolled", + )} + > + + + + { + navigate({ to: "/projects" }); + closeMobileIfOpen(); + }} + className="group/projects-item relative" + > + + + { + navigate({ to: "/hub" }); + closeMobileIfOpen(); + }} + /> + {/* Train has a labelled section when expanded; plain icon here only when collapsed. */} + { + if (chatOnly) return; + navigate({ to: "/studio" }); + closeMobileIfOpen(); + }} + className="hidden group-data-[collapsible=icon]:block" + /> + + + - { - navigate({ to: "/data-recipes" }); - closeMobileIfOpen(); - }} - /> - - { - if (chatOnly) return; - navigate({ to: "/export" }); - closeMobileIfOpen(); - }} - /> - - - - - - {/* Recent Chats — hide on Studio only (Eyera fac13); chatOpen = ec695 clickability */} - {!isStudioRoute && chatItems.length > 0 && ( - - - - - Recents - + + + + + {t("shell.navigation.train")} + - - - {chatItems.map((item) => ( - - { - navigate({ - to: "/chat", - search: - item.type === "single" - ? { thread: item.id } - : { compare: item.id }, - }); - closeMobileIfOpen(); - }} - > - {item.title} - - - - - - - openRenameChat(item)}> - - Rename - - setConfirmingDelete({ kind: "chat", item })} - > - - Delete - - - - - ))} - - + + + { + if (chatOnly) return; + navigate({ to: "/studio" }); + closeMobileIfOpen(); + }} + /> + { + navigate({ to: "/data-recipes" }); + closeMobileIfOpen(); + }} + /> + { + if (chatOnly) return; + navigate({ to: "/export" }); + closeMobileIfOpen(); + }} + /> + + + + + {!isStudioRoute && ( + + + +
+ + {t("shell.navigation.recents")} + + +
+ + + + + {recentChatItems.map((item) => + renderChatSidebarItem(item, "recent"), + )} + + + + )} - {/* Recent Runs */} {isStudioRoute && runItems.length > 0 && !chatOnly && ( - - - Recents - + + + {t("shell.navigation.recents")} + {runItems.map((run) => { + // Explicit selection wins. Otherwise highlight the active + // job only while the "Current Run" tab is the view, keeping + // the Configure tab unhighlighted even though activeJobId + // stays pinned to the last job. const isActiveRun = - selectedHistoryRunId === run.id || activeJobId === run.id; + selectedHistoryRunId != null + ? run.id === selectedHistoryRunId + : currentRunViewActive && run.id === activeJobId; return ( { setSelectedHistoryRunId(run.id); closeMobileIfOpen(); @@ -641,7 +1019,7 @@ export function AppSidebar() { @@ -841,10 +1272,14 @@ export function AppSidebar() { if (!open) setRenamingTarget(null); }} > - + - {renamingTarget?.kind === "run" ? "Rename run" : "Rename chat"} + {renamingTarget?.kind === "run" + ? t("shell.dialog.renameRun.title") + : renamingTarget?.kind === "project" + ? "Rename project" + : t("shell.dialog.renameChat.title")} @@ -868,14 +1315,66 @@ export function AppSidebar() { variant="ghost" onClick={() => setRenamingTarget(null)} > - Cancel + {t("common.cancel")} + + + + { + setCreatingProject(open); + if (!open) { + setProjectNameDraft(""); + setProjectCreateMoveTarget(null); + } + }} + > + + + + {projectCreateMoveTarget ? "Move to new project" : "New project"} + + + setProjectNameDraft(event.target.value)} + onKeyDown={(event) => { + if (event.key === "Enter") { + event.preventDefault(); + void commitCreateProject(); + } + }} + autoFocus + maxLength={120} + placeholder="Project name" + aria-label="Project name" + className="focus-visible:border-input focus-visible:ring-0" + /> + + + diff --git a/studio/frontend/src/components/assistant-ui/attachment.tsx b/studio/frontend/src/components/assistant-ui/attachment.tsx index 3ae1c68561..0a59d75964 100644 --- a/studio/frontend/src/components/assistant-ui/attachment.tsx +++ b/studio/frontend/src/components/assistant-ui/attachment.tsx @@ -24,7 +24,9 @@ import { useAui, useAuiState, } from "@assistant-ui/react"; -import { FileText, PlusIcon, XIcon } from "lucide-react"; +import { File02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { PlusIcon, XIcon } from "lucide-react"; import { type FC, type PropsWithChildren, @@ -134,7 +136,11 @@ const AttachmentThumb: FC = () => { return (
- +
); }; @@ -158,8 +164,8 @@ const AttachmentUI: FC = () => { throw new Error(`Unknown attachment type: ${type as string}`); } }); - // Include filename in accessible name so screen readers distinguish - // same-typed attachments. Sighted users get it via the tooltip. + // Filename in accessible name lets screen readers distinguish same-typed + // attachments. Sighted users get it via the tooltip. const accessibleName = name ? `${typeLabel} attachment: ${name}` : `${typeLabel} attachment`; @@ -238,7 +244,7 @@ export const ComposerAddAttachment: FC = () => { side="bottom" variant="ghost" size="icon" - className="aui-composer-add-attachment size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:border-muted-foreground/15 dark:hover:bg-muted-foreground/30" + className="aui-composer-add-attachment size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:hover:bg-muted-foreground/30" aria-label="Add Attachment" > diff --git a/studio/frontend/src/components/assistant-ui/audio-player.tsx b/studio/frontend/src/components/assistant-ui/audio-player.tsx index 6b99de8374..8836eddff5 100644 --- a/studio/frontend/src/components/assistant-ui/audio-player.tsx +++ b/studio/frontend/src/components/assistant-ui/audio-player.tsx @@ -4,7 +4,9 @@ "use client"; import { Button } from "@/components/ui/button"; -import { DownloadIcon, PauseIcon, PlayIcon } from "lucide-react"; +import { PauseIcon, PlayIcon } from "lucide-react"; +import { Download01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { type FC, useRef, useState } from "react"; interface AudioPlayerProps { @@ -110,7 +112,7 @@ export const AudioPlayer: FC = ({ src }) => { onClick={handleDownload} title="Download audio" > - +
); diff --git a/studio/frontend/src/components/assistant-ui/citation-utils.ts b/studio/frontend/src/components/assistant-ui/citation-utils.ts new file mode 100644 index 0000000000..58000c038f --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/citation-utils.ts @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +const RAG_SOURCES_SENTINEL = "__RAG_SOURCES__:"; + +export interface Citation { + id: string; + filename: string; + page?: number | null; + score?: number | null; + text: string; + documentId?: string | null; + chunkId?: string | null; +} + +function asNumber(value: unknown): number | null { + return typeof value === "number" && Number.isFinite(value) ? value : null; +} + +// null if absent so callers fall back to generic JSON shapes. +function parseSentinelSources(result: unknown): Citation[] | null { + if (typeof result !== "string") return null; + const idx = result.indexOf(RAG_SOURCES_SENTINEL); + if (idx < 0) return null; + const payload = result.slice(idx + RAG_SOURCES_SENTINEL.length).trim(); + let rows: unknown; + try { + rows = JSON.parse(payload); + } catch { + return []; + } + if (!Array.isArray(rows)) return []; + return rows.map((row, i) => { + const r = (row ?? {}) as Record; + const documentId = typeof r.documentId === "string" ? r.documentId : null; + const chunkId = typeof r.chunkId === "string" ? r.chunkId : null; + const filename = + typeof r.filename === "string" ? r.filename : `Source ${i + 1}`; + return { + id: chunkId ?? `${filename}-${i}`, + filename, + page: asNumber(r.page), + score: asNumber(r.score), + text: typeof r.text === "string" ? r.text : "", + documentId, + chunkId, + }; + }); +} + +export function parseCitations(result: unknown): Citation[] { + const sentinel = parseSentinelSources(result); + if (sentinel !== null) return sentinel; + + let rows: unknown[] | null = null; + if (Array.isArray(result)) { + rows = result; + } else if (typeof result === "string") { + const trimmed = result.trim(); + if (trimmed.startsWith("[") || trimmed.startsWith("{")) { + try { + const parsed = JSON.parse(trimmed); + if (Array.isArray(parsed)) rows = parsed; + else if (parsed && Array.isArray((parsed as { results?: unknown }).results)) { + rows = (parsed as { results: unknown[] }).results; + } + } catch { + rows = null; + } + } + } else if (result && Array.isArray((result as { results?: unknown }).results)) { + rows = (result as { results: unknown[] }).results; + } + if (!rows) return []; + + const citations: Citation[] = []; + rows.forEach((row, i) => { + if (!row || typeof row !== "object") return; + const r = row as Record; + const text = + typeof r.text === "string" + ? r.text + : typeof r.chunk === "string" + ? r.chunk + : typeof r.content === "string" + ? r.content + : ""; + const filename = + typeof r.filename === "string" + ? r.filename + : typeof r.documentId === "string" + ? r.documentId + : `Source ${i + 1}`; + const chunkId = + typeof r.chunkId === "string" ? r.chunkId : `${filename}-${i}`; + citations.push({ + id: chunkId, + filename, + page: asNumber(r.page), + score: asNumber(r.score), + text, + }); + }); + return citations; +} diff --git a/studio/frontend/src/components/assistant-ui/code-plugin.ts b/studio/frontend/src/components/assistant-ui/code-plugin.ts index 5df7ac4f95..1e70871c06 100644 --- a/studio/frontend/src/components/assistant-ui/code-plugin.ts +++ b/studio/frontend/src/components/assistant-ui/code-plugin.ts @@ -10,8 +10,8 @@ import { } from "@streamdown/code"; import type { BundledLanguage } from "shiki"; -// Fence tags LLMs/users commonly write that shiki doesn't expose as aliases. -// Keys are lower-cased input; values are canonical shiki language ids. +// Common fence tags shiki doesn't expose as aliases. +// Keys: lower-cased input; values: canonical shiki language ids. const LANGUAGE_ALIAS_OVERRIDES: Record = { objectivec: "objective-c", "obj-c": "objective-c", diff --git a/studio/frontend/src/components/assistant-ui/code-themes.ts b/studio/frontend/src/components/assistant-ui/code-themes.ts index 2557b45ef0..67db02b9b4 100644 --- a/studio/frontend/src/components/assistant-ui/code-themes.ts +++ b/studio/frontend/src/components/assistant-ui/code-themes.ts @@ -5,11 +5,10 @@ import oneDarkPro from "@shikijs/themes/one-dark-pro"; import oneLight from "@shikijs/themes/one-light"; import type { ThemeRegistrationAny } from "shiki"; -// Canonical Atom One Dark / One Light themes, shipped by `@shikijs/themes`. -// We only override the background so the code block blends into the app's -// `--code-block` surface instead of painting its own. Every token color and -// scope mapping is left intact — that's what gives consistent multi-language -// highlighting (including Objective-C, Go, Rust, etc.) out of the box. +// Canonical Atom One Dark / One Light themes from `@shikijs/themes`. Only the +// background is overridden so the code block blends into the app's `--code-block` +// surface; all token colors/scopes are kept intact for consistent multi-language +// highlighting out of the box. const withTransparentBg = (theme: ThemeRegistrationAny): ThemeRegistrationAny => ({ ...theme, bg: "transparent", diff --git a/studio/frontend/src/components/assistant-ui/image.tsx b/studio/frontend/src/components/assistant-ui/image.tsx index a2e3f30dc1..0850a8230a 100644 --- a/studio/frontend/src/components/assistant-ui/image.tsx +++ b/studio/frontend/src/components/assistant-ui/image.tsx @@ -15,13 +15,14 @@ import type { import { type VariantProps, cva } from "class-variance-authority"; import { CopyIcon, - DownloadIcon, ImageIcon, ImageOffIcon, Loader2Icon, RefreshCwIcon, ShieldAlertIcon, } from "lucide-react"; +import { Download01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { type ComponentProps, type PropsWithChildren, @@ -377,10 +378,7 @@ function ImageContentFilterError({ export type ImageActionsProps = { part: ImageMessagePart; - /** - * Wire to your own generation call to show a regenerate button. The button - * renders only when this is set and the part carries a `prompt`. - */ + /** Shows a regenerate button (only when set and the part has a `prompt`). */ onRegenerate?: () => void | Promise; className?: string; }; @@ -427,7 +425,7 @@ function ImageActions({ part, onRegenerate, className }: ImageActionsProps) { aria-label="Download image" className="inline-flex size-7 items-center justify-center rounded hover:bg-muted" > - + -
-
-